diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 8253b52..ce62ccf 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -12616,6 +12616,7 @@
     field public static final String ACTION_REQUEST_ENABLE_CONTENT_CAPTURE = "android.settings.REQUEST_ENABLE_CONTENT_CAPTURE";
     field public static final String ACTION_SHOW_ADMIN_SUPPORT_DETAILS = "android.settings.SHOW_ADMIN_SUPPORT_DETAILS";
     field public static final String ACTION_SHOW_RESTRICTED_SETTING_DIALOG = "android.settings.SHOW_RESTRICTED_SETTING_DIALOG";
+    field @FlaggedApi("com.android.internal.telephony.flags.action_sim_preference_settings") public static final String ACTION_SIM_PREFERENCE_SETTINGS = "android.settings.SIM_PREFERENCE_SETTINGS";
     field public static final String ACTION_TETHER_PROVISIONING_UI = "android.settings.TETHER_PROVISIONING_UI";
     field public static final String ACTION_TETHER_SETTINGS = "android.settings.TETHER_SETTINGS";
     field public static final String ACTION_TETHER_UNSUPPORTED_CARRIER_UI = "android.settings.TETHER_UNSUPPORTED_CARRIER_UI";
diff --git a/core/java/android/os/OWNERS b/core/java/android/os/OWNERS
index bfcc5cc..8d35338 100644
--- a/core/java/android/os/OWNERS
+++ b/core/java/android/os/OWNERS
@@ -118,9 +118,10 @@
 # Memory
 per-file OomKillRecord.java = file:/MEMORY_OWNERS
 
-# MessageQueue
+# MessageQueue and related classes
 per-file MessageQueue.java = mfasheh@google.com, shayba@google.com
 per-file Message.java = mfasheh@google.com, shayba@google.com
+per-file TestLooperManager.java = mfasheh@google.com, shayba@google.com
 
 # Stats
 per-file IStatsBootstrapAtomService.aidl = file:/services/core/java/com/android/server/stats/OWNERS
diff --git a/core/java/android/provider/Settings.java b/core/java/android/provider/Settings.java
index 0ae9ffa..9935be2 100644
--- a/core/java/android/provider/Settings.java
+++ b/core/java/android/provider/Settings.java
@@ -2469,6 +2469,25 @@
             = "android.settings.APP_NOTIFICATION_BUBBLE_SETTINGS";
 
     /**
+     * Activity Action: Show the settings for users to select their preferred SIM subscription
+     * when a new SIM subscription has become available.
+     * <p>
+     * This Activity will only launch successfully if the newly active subscription ID is set as the
+     * value of {@link EXTRA_SUB_ID} and the value corresponds with an active SIM subscription.
+     * <p>
+     * Input: {@link #EXTRA_SUB_ID}: the subscription ID of the newly active SIM subscription.
+     * <p>
+     * Output: Nothing.
+     *
+     * @hide
+     */
+    @FlaggedApi(com.android.internal.telephony.flags.Flags.FLAG_ACTION_SIM_PREFERENCE_SETTINGS)
+    @SystemApi
+    @SdkConstant(SdkConstantType.ACTIVITY_INTENT_ACTION)
+    public static final String ACTION_SIM_PREFERENCE_SETTINGS =
+            "android.settings.SIM_PREFERENCE_SETTINGS";
+
+    /**
      * Intent Extra: The value of {@link android.app.settings.SettingsEnums#EntryPointType} for
      * settings metrics that logs the entry point about physical keyboard settings.
      * <p>
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index 596a9af..13c125c 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -2372,7 +2372,7 @@
     <string name="default_sms_application" translatable="false">com.android.messaging</string>
 
     <!-- Flag indicating whether the current device supports "Ask every time" for sms-->
-    <bool name="config_sms_ask_every_time_support">true</bool>
+    <bool name="config_sms_ask_every_time_support">false</bool>
 
     <!-- Flag indicating whether the current device allows acknowledgement of SIM operation like
          SM-PP or saving SMS to SIM can be done via the IMS interfaces.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index f46b955..86e0d08 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -152,6 +152,7 @@
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalSystemViewContainer;
 import com.android.wm.shell.windowdecor.common.viewhost.DefaultWindowDecorViewHostSupplier;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
 import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.education.DesktopWindowingEducationPromoController;
 import com.android.wm.shell.windowdecor.education.DesktopWindowingEducationTooltipController;
@@ -301,6 +302,7 @@
             Transitions transitions,
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer,
             FocusTransitionObserver focusTransitionObserver,
+            WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             Optional<DesktopModeWindowDecorViewModel> desktopModeWindowDecorViewModel) {
         if (desktopModeWindowDecorViewModel.isPresent()) {
             return desktopModeWindowDecorViewModel.get();
@@ -318,7 +320,8 @@
                 rootTaskDisplayAreaOrganizer,
                 syncQueue,
                 transitions,
-                focusTransitionObserver);
+                focusTransitionObserver,
+                windowDecorViewHostSupplier);
     }
 
     @WMSingleton
@@ -343,7 +346,7 @@
 
     @WMSingleton
     @Provides
-    static WindowDecorViewHostSupplier provideWindowDecorViewHostSupplier(
+    static WindowDecorViewHostSupplier<WindowDecorViewHost> provideWindowDecorViewHostSupplier(
             @ShellMainThread @NonNull CoroutineScope mainScope) {
         return new DefaultWindowDecorViewHostSupplier(mainScope);
     }
@@ -908,6 +911,7 @@
             InteractionJankMonitor interactionJankMonitor,
             AppToWebGenericLinksParser genericLinksParser,
             AssistContentRequester assistContentRequester,
+            WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             MultiInstanceHelper multiInstanceHelper,
             Optional<DesktopTasksLimiter> desktopTasksLimiter,
             AppHandleEducationController appHandleEducationController,
@@ -927,8 +931,8 @@
                 displayInsetsController, syncQueue, transitions, desktopTasksController,
                 desktopImmersiveController.get(),
                 rootTaskDisplayAreaOrganizer, interactionJankMonitor, genericLinksParser,
-                assistContentRequester, multiInstanceHelper, desktopTasksLimiter,
-                appHandleEducationController, appToWebEducationController,
+                assistContentRequester, windowDecorViewHostSupplier, multiInstanceHelper,
+                desktopTasksLimiter, appHandleEducationController, appToWebEducationController,
                 windowDecorCaptionHandleRepository, activityOrientationChangeHandler,
                 focusTransitionObserver, desktopModeEventLogger, desktopModeUiEventLogger));
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
index 885f3db..0b91966 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
@@ -68,6 +68,8 @@
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.FocusTransitionObserver;
 import com.android.wm.shell.transition.Transitions;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.extension.TaskInfoKt;
 
 /**
@@ -90,6 +92,7 @@
     private final Transitions mTransitions;
     private final Region mExclusionRegion = Region.obtain();
     private final InputManager mInputManager;
+    private final WindowDecorViewHostSupplier<WindowDecorViewHost> mWindowDecorViewHostSupplier;
     private TaskOperations mTaskOperations;
     private FocusTransitionObserver mFocusTransitionObserver;
 
@@ -130,7 +133,8 @@
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer,
             SyncTransactionQueue syncQueue,
             Transitions transitions,
-            FocusTransitionObserver focusTransitionObserver) {
+            FocusTransitionObserver focusTransitionObserver,
+            WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier) {
         mContext = context;
         mMainExecutor = shellExecutor;
         mMainHandler = mainHandler;
@@ -143,6 +147,7 @@
         mSyncQueue = syncQueue;
         mTransitions = transitions;
         mFocusTransitionObserver = focusTransitionObserver;
+        mWindowDecorViewHostSupplier = windowDecorViewHostSupplier;
         if (!Transitions.ENABLE_SHELL_TRANSITIONS) {
             mTaskOperations = new TaskOperations(null, mContext, mSyncQueue);
         }
@@ -332,7 +337,8 @@
                         mMainHandler,
                         mBgExecutor,
                         mMainChoreographer,
-                        mSyncQueue);
+                        mSyncQueue,
+                        mWindowDecorViewHostSupplier);
         mWindowDecorByTaskId.put(taskInfo.taskId, windowDecoration);
 
         final FluidResizeTaskPositioner taskPositioner =
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
index 112e429..23bb2aa 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
@@ -58,6 +58,8 @@
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.extension.TaskInfoKt;
 
 /**
@@ -90,9 +92,10 @@
             Handler handler,
             @ShellBackgroundThread ShellExecutor bgExecutor,
             Choreographer choreographer,
-            SyncTransactionQueue syncQueue) {
-        super(context, userContext, displayController, taskOrganizer, handler, taskInfo,
-                taskSurface);
+            SyncTransactionQueue syncQueue,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier) {
+        super(context, userContext, displayController, taskOrganizer, taskInfo,
+                taskSurface, windowDecorViewHostSupplier);
         mHandler = handler;
         mBgExecutor = bgExecutor;
         mChoreographer = choreographer;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index e8b02dc..08d047a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -139,6 +139,8 @@
 import com.android.wm.shell.transition.FocusTransitionObserver;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.DesktopModeWindowDecoration.ExclusionRegionListener;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.extension.InsetsStateKt;
 import com.android.wm.shell.windowdecor.extension.TaskInfoKt;
 import com.android.wm.shell.windowdecor.viewholder.AppHeaderViewHolder;
@@ -217,6 +219,7 @@
     private boolean mInImmersiveMode;
     private final String mSysUIPackageName;
     private final AssistContentRequester mAssistContentRequester;
+    private final WindowDecorViewHostSupplier<WindowDecorViewHost> mWindowDecorViewHostSupplier;
 
     private final DisplayChangeController.OnDisplayChangingListener mOnDisplayChangingListener;
     private final ISystemGestureExclusionListener mGestureExclusionListener =
@@ -260,6 +263,7 @@
             InteractionJankMonitor interactionJankMonitor,
             AppToWebGenericLinksParser genericLinksParser,
             AssistContentRequester assistContentRequester,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             MultiInstanceHelper multiInstanceHelper,
             Optional<DesktopTasksLimiter> desktopTasksLimiter,
             AppHandleEducationController appHandleEducationController,
@@ -289,6 +293,7 @@
                 desktopImmersiveController,
                 genericLinksParser,
                 assistContentRequester,
+                windowDecorViewHostSupplier,
                 multiInstanceHelper,
                 new DesktopModeWindowDecoration.Factory(),
                 new InputMonitorFactory(),
@@ -329,6 +334,7 @@
             DesktopImmersiveController desktopImmersiveController,
             AppToWebGenericLinksParser genericLinksParser,
             AssistContentRequester assistContentRequester,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             MultiInstanceHelper multiInstanceHelper,
             DesktopModeWindowDecoration.Factory desktopModeWindowDecorFactory,
             InputMonitorFactory inputMonitorFactory,
@@ -381,6 +387,7 @@
         mWindowDecorCaptionHandleRepository = windowDecorCaptionHandleRepository;
         mActivityOrientationChangeHandler = activityOrientationChangeHandler;
         mAssistContentRequester = assistContentRequester;
+        mWindowDecorViewHostSupplier = windowDecorViewHostSupplier;
         mOnDisplayChangingListener = (displayId, fromRotation, toRotation, displayAreaInfo, t) -> {
             DesktopModeWindowDecoration decoration;
             RunningTaskInfo taskInfo;
@@ -1623,6 +1630,7 @@
                         mRootTaskDisplayAreaOrganizer,
                         mGenericLinksParser,
                         mAssistContentRequester,
+                        mWindowDecorViewHostSupplier,
                         mMultiInstanceHelper,
                         mWindowDecorCaptionHandleRepository,
                         mDesktopModeEventLogger);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index 1263176..6562f38 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -105,6 +105,8 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource;
 import com.android.wm.shell.shared.multiinstance.ManageWindowsViewContainer;
 import com.android.wm.shell.splitscreen.SplitScreenController;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.extension.TaskInfoKt;
 import com.android.wm.shell.windowdecor.viewholder.AppHandleViewHolder;
 import com.android.wm.shell.windowdecor.viewholder.AppHeaderViewHolder;
@@ -219,6 +221,7 @@
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer,
             AppToWebGenericLinksParser genericLinksParser,
             AssistContentRequester assistContentRequester,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             MultiInstanceHelper multiInstanceHelper,
             WindowDecorCaptionHandleRepository windowDecorCaptionHandleRepository,
             DesktopModeEventLogger desktopModeEventLogger) {
@@ -230,6 +233,7 @@
                 WindowContainerTransaction::new, SurfaceControl::new, new WindowManagerWrapper(
                         context.getSystemService(WindowManager.class)),
                 new SurfaceControlViewHostFactory() {},
+                windowDecorViewHostSupplier,
                 DefaultMaximizeMenuFactory.INSTANCE,
                 DefaultHandleMenuFactory.INSTANCE, multiInstanceHelper,
                 windowDecorCaptionHandleRepository, desktopModeEventLogger);
@@ -258,15 +262,16 @@
             Supplier<SurfaceControl> surfaceControlSupplier,
             WindowManagerWrapper windowManagerWrapper,
             SurfaceControlViewHostFactory surfaceControlViewHostFactory,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             MaximizeMenuFactory maximizeMenuFactory,
             HandleMenuFactory handleMenuFactory,
             MultiInstanceHelper multiInstanceHelper,
             WindowDecorCaptionHandleRepository windowDecorCaptionHandleRepository,
             DesktopModeEventLogger desktopModeEventLogger) {
-        super(context, userContext, displayController, taskOrganizer, handler, taskInfo,
+        super(context, userContext, displayController, taskOrganizer, taskInfo,
                 taskSurface, surfaceControlBuilderSupplier, surfaceControlTransactionSupplier,
                 windowContainerTransactionSupplier, surfaceControlSupplier,
-                surfaceControlViewHostFactory, desktopModeEventLogger);
+                surfaceControlViewHostFactory, windowDecorViewHostSupplier, desktopModeEventLogger);
         mSplitScreenController = splitScreenController;
         mHandler = handler;
         mBgExecutor = bgExecutor;
@@ -1760,6 +1765,8 @@
                 RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer,
                 AppToWebGenericLinksParser genericLinksParser,
                 AssistContentRequester assistContentRequester,
+                @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost>
+                        windowDecorViewHostSupplier,
                 MultiInstanceHelper multiInstanceHelper,
                 WindowDecorCaptionHandleRepository windowDecorCaptionHandleRepository,
                 DesktopModeEventLogger desktopModeEventLogger) {
@@ -1780,6 +1787,7 @@
                     rootTaskDisplayAreaOrganizer,
                     genericLinksParser,
                     assistContentRequester,
+                    windowDecorViewHostSupplier,
                     multiInstanceHelper,
                     windowDecorCaptionHandleRepository,
                     desktopModeEventLogger);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
index 261d400..5d1bedb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
@@ -39,7 +39,6 @@
 import android.graphics.Rect;
 import android.graphics.Region;
 import android.os.Binder;
-import android.os.Handler;
 import android.os.Trace;
 import android.view.Display;
 import android.view.InsetsSource;
@@ -59,10 +58,11 @@
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.desktopmode.DesktopModeEventLogger;
-import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.windowdecor.WindowDecoration.RelayoutParams.OccludingCaptionElement;
 import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewHostViewContainer;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.extension.InsetsStateKt;
 
 import java.util.ArrayList;
@@ -90,10 +90,10 @@
         implements AutoCloseable {
 
     /**
-     * The Z-order of {@link #mCaptionContainerSurface}.
+     * The Z-order of the caption surface.
      * <p>
      * We use {@link #mDecorationContainerSurface} to define input window for task resizing; by
-     * layering it in front of {@link #mCaptionContainerSurface}, we can allow it to handle input
+     * layering it in front of the caption surface, we can allow it to handle input
      * prior to caption view itself, treating corner inputs as resize events rather than
      * repositioning.
      */
@@ -102,7 +102,7 @@
      * The Z-order of the task input sink in {@link DragPositioningCallback}.
      * <p>
      * This task input sink is used to prevent undesired dispatching of motion events out of task
-     * bounds; by layering it behind {@link #mCaptionContainerSurface}, we allow captions to handle
+     * bounds; by layering it behind the caption surface, we allow captions to handle
      * input events first.
      */
     static final int INPUT_SINK_Z_ORDER = -2;
@@ -123,12 +123,13 @@
     final @NonNull Context mUserContext;
     final @NonNull DisplayController mDisplayController;
     final @NonNull DesktopModeEventLogger mDesktopModeEventLogger;
-    private final @ShellMainThread Handler mMainHandler;
     final ShellTaskOrganizer mTaskOrganizer;
     final Supplier<SurfaceControl.Builder> mSurfaceControlBuilderSupplier;
     final Supplier<SurfaceControl.Transaction> mSurfaceControlTransactionSupplier;
     final Supplier<WindowContainerTransaction> mWindowContainerTransactionSupplier;
     final SurfaceControlViewHostFactory mSurfaceControlViewHostFactory;
+    @NonNull private final WindowDecorViewHostSupplier<WindowDecorViewHost>
+            mWindowDecorViewHostSupplier;
     private final DisplayController.OnDisplaysChangedListener mOnDisplaysChangedListener =
             new DisplayController.OnDisplaysChangedListener() {
                 @Override
@@ -153,9 +154,7 @@
     Display mDisplay;
     SurfaceControl mDecorationContainerSurface;
 
-    SurfaceControl mCaptionContainerSurface;
-    private CaptionWindowlessWindowManager mCaptionWindowManager;
-    private SurfaceControlViewHost mViewHost;
+    private WindowDecorViewHost mViewHost;
     private Configuration mWindowDecorConfig;
     TaskDragResizer mTaskDragResizer;
     boolean mIsCaptionVisible;
@@ -170,20 +169,19 @@
     private final Binder mOwner = new Binder();
     private final float[] mTmpColor = new float[3];
 
-    private Runnable mCurrentUpdateViewHostRunnable;
-
     WindowDecoration(
             Context context,
             @NonNull Context userContext,
             DisplayController displayController,
             ShellTaskOrganizer taskOrganizer,
-            @ShellMainThread Handler mainHandler,
             RunningTaskInfo taskInfo,
-            SurfaceControl taskSurface) {
-        this(context, userContext, displayController, taskOrganizer, mainHandler, taskInfo,
+            SurfaceControl taskSurface,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier) {
+        this(context, userContext, displayController, taskOrganizer, taskInfo,
                 taskSurface, SurfaceControl.Builder::new, SurfaceControl.Transaction::new,
                 WindowContainerTransaction::new, SurfaceControl::new,
-                new SurfaceControlViewHostFactory() {}, new DesktopModeEventLogger());
+                new SurfaceControlViewHostFactory() {}, windowDecorViewHostSupplier,
+                new DesktopModeEventLogger());
     }
 
     WindowDecoration(
@@ -191,7 +189,6 @@
             @NonNull Context userContext,
             @NonNull DisplayController displayController,
             ShellTaskOrganizer taskOrganizer,
-            @ShellMainThread Handler mainHandler,
             RunningTaskInfo taskInfo,
             @NonNull SurfaceControl taskSurface,
             Supplier<SurfaceControl.Builder> surfaceControlBuilderSupplier,
@@ -199,13 +196,13 @@
             Supplier<WindowContainerTransaction> windowContainerTransactionSupplier,
             Supplier<SurfaceControl> surfaceControlSupplier,
             SurfaceControlViewHostFactory surfaceControlViewHostFactory,
+            @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost> windowDecorViewHostSupplier,
             @NonNull DesktopModeEventLogger desktopModeEventLogger
     ) {
         mContext = context;
         mUserContext = userContext;
         mDisplayController = displayController;
         mTaskOrganizer = taskOrganizer;
-        mMainHandler = mainHandler;
         mTaskInfo = taskInfo;
         mTaskSurface = cloneSurfaceControl(taskSurface, surfaceControlSupplier);
         mDesktopModeEventLogger = desktopModeEventLogger;
@@ -213,6 +210,7 @@
         mSurfaceControlTransactionSupplier = surfaceControlTransactionSupplier;
         mWindowContainerTransactionSupplier = windowContainerTransactionSupplier;
         mSurfaceControlViewHostFactory = surfaceControlViewHostFactory;
+        mWindowDecorViewHostSupplier = windowDecorViewHostSupplier;
         mDisplay = mDisplayController.getDisplay(mTaskInfo.displayId);
         final InsetsState insetsState = mDisplayController.getInsetsState(mTaskInfo.displayId);
         mIsStatusBarVisible = insetsState != null
@@ -292,43 +290,20 @@
         outResult.mCaptionY = 0;
         outResult.mCaptionTopPadding = params.mCaptionTopPadding;
 
+        Trace.beginSection("relayout-createViewHostIfNeeded");
+        createViewHostIfNeeded(mDecorWindowContext, mDisplay);
+        Trace.endSection();
+
         Trace.beginSection("WindowDecoration#relayout-updateSurfacesAndInsets");
+        final SurfaceControl captionSurface = mViewHost.getSurfaceControl();
         updateDecorationContainerSurface(startT, outResult);
-        final SurfaceControl captionSurface = getOrCreateCaptionSurface();
         updateCaptionContainerSurface(captionSurface, startT, outResult);
         updateCaptionInsets(params, wct, outResult, taskBounds);
         updateTaskSurface(params, startT, finishT, outResult);
         Trace.endSection();
 
         Trace.beginSection("WindowDecoration#relayout-updateViewHost");
-        clearCurrentViewHostRunnable();
-        if (params.mAsyncViewHost) {
-            mCurrentUpdateViewHostRunnable = () -> updateViewHost(params, startT, outResult);
-            mMainHandler.post(mCurrentUpdateViewHostRunnable);
-        } else {
-            updateViewHost(params, startT, outResult);
-        }
-        Trace.endSection();
-
-        Trace.endSection(); // WindowDecoration#relayout
-    }
-
-    private void clearCurrentViewHostRunnable() {
-        if (mCurrentUpdateViewHostRunnable != null) {
-            mMainHandler.removeCallbacks(mCurrentUpdateViewHostRunnable);
-            mCurrentUpdateViewHostRunnable = null;
-        }
-    }
-
-    private void updateViewHost(RelayoutParams params, SurfaceControl.Transaction startT,
-            RelayoutResult<T> outResult) {
-        createCaptionWindowManagerIfNeeded();
-        Trace.beginSection("updateViewHost-createViewHostIfNeeded");
-        final boolean firstLayout = createViewHostIfNeeded(mDecorWindowContext, mDisplay);
-        Trace.endSection();
-
         outResult.mRootView.setPadding(0, params.mCaptionTopPadding, 0, 0);
-        mCaptionWindowManager.setConfiguration(mTaskInfo.getConfiguration());
         final Rect localCaptionBounds = new Rect(
                 outResult.mCaptionX,
                 outResult.mCaptionY,
@@ -337,50 +312,21 @@
         final Region touchableRegion = params.mLimitTouchRegionToSystemAreas
                 ? calculateLimitedTouchableRegion(params, localCaptionBounds)
                 : null;
-        if (params.mLimitTouchRegionToSystemAreas) {
-            mCaptionWindowManager.setTouchRegion(mViewHost, touchableRegion);
-        }
-        if (touchableRegion != null) {
-            touchableRegion.recycle();
-        }
-        updateViewHierarchy(params, outResult, startT, firstLayout);
+        updateViewHierarchy(params, outResult, startT, touchableRegion);
+        Trace.endSection();
+
+        Trace.endSection(); // WindowDecoration#relayout
     }
 
-    private SurfaceControl getOrCreateCaptionSurface() {
-        if (mCaptionContainerSurface == null) {
-            final SurfaceControl.Builder builder = mSurfaceControlBuilderSupplier.get();
-            mCaptionContainerSurface = builder
-                    .setName("Caption container of Task=" + mTaskInfo.taskId)
-                    .setContainerLayer()
-                    .setParent(mDecorationContainerSurface)
-                    .setCallsite("WindowDecoration.updateCaptionContainerSurface")
-                    .build();
-        }
-        return mCaptionContainerSurface;
-    }
-
-    private boolean createViewHostIfNeeded(@NonNull Context context, @NonNull Display display) {
+    private void createViewHostIfNeeded(@NonNull Context context, @NonNull Display display) {
         if (mViewHost == null) {
-            mViewHost = mSurfaceControlViewHostFactory.create(context, display,
-                    mCaptionWindowManager);
-            Trace.endSection();
-            return true;
-        }
-        return false;
-    }
-
-    private void createCaptionWindowManagerIfNeeded() {
-        if (mCaptionWindowManager == null) {
-            // Put caption under a container surface because ViewRootImpl sets the destination frame
-            // of windowless window layers and BLASTBufferQueue#update() doesn't support offset.
-            mCaptionWindowManager = new CaptionWindowlessWindowManager(
-                    mTaskInfo.getConfiguration(), mCaptionContainerSurface);
+            mViewHost = mWindowDecorViewHostSupplier.acquire(context, display);
         }
     }
 
     private void updateViewHierarchy(@NonNull RelayoutParams params,
             @NonNull RelayoutResult<T> outResult, @NonNull SurfaceControl.Transaction startT,
-            boolean firstLayout) {
+            @Nullable Region touchableRegion) {
         Trace.beginSection("WindowDecoration#updateViewHierarchy");
         final WindowManager.LayoutParams lp =
                 new WindowManager.LayoutParams(
@@ -392,16 +338,15 @@
         lp.setTitle("Caption of Task=" + mTaskInfo.taskId);
         lp.setTrustedOverlay();
         lp.inputFeatures = params.mInputFeatures;
-        if (params.mApplyStartTransactionOnDraw) {
-            if (params.mAsyncViewHost) {
+        if (params.mAsyncViewHost) {
+            if (params.mApplyStartTransactionOnDraw) {
                 throw new IllegalArgumentException("Cannot use sync draw tx with async relayout");
             }
-            mViewHost.getRootSurfaceControl().applyTransactionOnDraw(startT);
-        }
-        if (firstLayout) {
-            mViewHost.setView(outResult.mRootView, lp);
+            mViewHost.updateViewAsync(outResult.mRootView, lp, mTaskInfo.configuration,
+                    touchableRegion);
         } else {
-            mViewHost.relayout(lp);
+            mViewHost.updateView(outResult.mRootView, lp, mTaskInfo.configuration,
+                    touchableRegion, params.mApplyStartTransactionOnDraw ? startT : null);
         }
         Trace.endSection();
     }
@@ -719,18 +664,11 @@
     }
 
     void releaseViews(WindowContainerTransaction wct) {
-        if (mViewHost != null) {
-            mViewHost.release();
-            mViewHost = null;
-        }
-
-        mCaptionWindowManager = null;
-
         final SurfaceControl.Transaction t = mSurfaceControlTransactionSupplier.get();
         boolean released = false;
-        if (mCaptionContainerSurface != null) {
-            t.remove(mCaptionContainerSurface);
-            mCaptionContainerSurface = null;
+        if (mViewHost != null) {
+            mWindowDecorViewHostSupplier.release(mViewHost, t);
+            mViewHost = null;
             released = true;
         }
 
@@ -753,7 +691,6 @@
     @Override
     public void close() {
         Trace.beginSection("WindowDecoration#close");
-        clearCurrentViewHostRunnable();
         mDisplayController.removeDisplayWindowListener(mOnDisplaysChangedListener);
         final WindowContainerTransaction wct = mWindowContainerTransactionSupplier.get();
         releaseViews(wct);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHost.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHost.kt
index c470eef..50aa21e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHost.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHost.kt
@@ -17,6 +17,7 @@
 
 import android.content.Context
 import android.content.res.Configuration
+import android.graphics.Region
 import android.view.Display
 import android.view.SurfaceControl
 import android.view.SurfaceControlViewHost
@@ -59,7 +60,7 @@
             .setCallsite("DefaultWindowDecorViewHost#init")
             .build()
 
-    private var wwm: WindowlessWindowManager? = null
+    private var wwm: WindowDecorWindowlessWindowManager? = null
     @VisibleForTesting var viewHost: SurfaceControlViewHost? = null
     private var currentUpdateJob: Job? = null
 
@@ -70,11 +71,12 @@
         view: View,
         attrs: WindowManager.LayoutParams,
         configuration: Configuration,
+        touchableRegion: Region?,
         onDrawTransaction: SurfaceControl.Transaction?,
     ) {
         Trace.beginSection("DefaultWindowDecorViewHost#updateView")
         clearCurrentUpdateJob()
-        updateViewHost(view, attrs, configuration, onDrawTransaction)
+        updateViewHost(view, attrs, configuration, touchableRegion, onDrawTransaction)
         Trace.endSection()
     }
 
@@ -82,12 +84,19 @@
         view: View,
         attrs: WindowManager.LayoutParams,
         configuration: Configuration,
+        touchableRegion: Region?,
     ) {
         Trace.beginSection("DefaultWindowDecorViewHost#updateViewAsync")
         clearCurrentUpdateJob()
         currentUpdateJob =
             mainScope.launch {
-                updateViewHost(view, attrs, configuration, onDrawTransaction = null)
+                updateViewHost(
+                    view,
+                    attrs,
+                    configuration,
+                    touchableRegion,
+                    onDrawTransaction = null
+                )
             }
         Trace.endSection()
     }
@@ -102,13 +111,13 @@
         view: View,
         attrs: WindowManager.LayoutParams,
         configuration: Configuration,
+        touchableRegion: Region?,
         onDrawTransaction: SurfaceControl.Transaction?,
     ) {
         Trace.beginSection("DefaultWindowDecorViewHost#updateViewHost")
         if (wwm == null) {
-            wwm = WindowlessWindowManager(configuration, rootSurface, null)
+            wwm = WindowDecorWindowlessWindowManager(configuration, rootSurface)
         }
-        requireWindowlessWindowManager().setConfiguration(configuration)
         if (viewHost == null) {
             viewHost =
                 surfaceControlViewHostFactory.invoke(
@@ -118,6 +127,10 @@
                     "DefaultWindowDecorViewHost#updateViewHost",
                 )
         }
+        requireWindowlessWindowManager().apply {
+            setConfiguration(configuration)
+            setTouchRegion(requireViewHost(), touchableRegion)
+        }
         onDrawTransaction?.let { requireViewHost().rootSurfaceControl.applyTransactionOnDraw(it) }
         if (requireViewHost().view == null) {
             Trace.beginSection("DefaultWindowDecorViewHost#updateViewHost-setView")
@@ -137,7 +150,7 @@
         currentUpdateJob = null
     }
 
-    private fun requireWindowlessWindowManager(): WindowlessWindowManager {
+    private fun requireWindowlessWindowManager(): WindowDecorWindowlessWindowManager {
         return wwm ?: error("Expected non-null windowless window manager")
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHostSupplier.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHostSupplier.kt
index 27ffd6c..7821619 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHostSupplier.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/DefaultWindowDecorViewHostSupplier.kt
@@ -24,14 +24,15 @@
 /**
  * A supplier of [DefaultWindowDecorViewHost]s. It creates a new one every time one is requested.
  */
-class DefaultWindowDecorViewHostSupplier(@ShellMainThread private val mainScope: CoroutineScope) :
-    WindowDecorViewHostSupplier<DefaultWindowDecorViewHost> {
+class DefaultWindowDecorViewHostSupplier(
+    @ShellMainThread private val mainScope: CoroutineScope
+) : WindowDecorViewHostSupplier<WindowDecorViewHost> {
 
-    override fun acquire(context: Context, display: Display): DefaultWindowDecorViewHost {
+    override fun acquire(context: Context, display: Display): WindowDecorViewHost {
         return DefaultWindowDecorViewHost(context, mainScope, display)
     }
 
-    override fun release(viewHost: DefaultWindowDecorViewHost, t: SurfaceControl.Transaction) {
+    override fun release(viewHost: WindowDecorViewHost, t: SurfaceControl.Transaction) {
         viewHost.release(t)
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorViewHost.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorViewHost.kt
index 7c1479e..2dcbbac 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorViewHost.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorViewHost.kt
@@ -16,6 +16,7 @@
 package com.android.wm.shell.windowdecor.common.viewhost
 
 import android.content.res.Configuration
+import android.graphics.Region
 import android.view.SurfaceControl
 import android.view.View
 import android.view.WindowManager
@@ -34,11 +35,17 @@
         view: View,
         attrs: WindowManager.LayoutParams,
         configuration: Configuration,
-        onDrawTransaction: SurfaceControl.Transaction?,
+        touchableRegion: Region? = null,
+        onDrawTransaction: SurfaceControl.Transaction? = null,
     )
 
     /** Asynchronously update the view hierarchy of this view host. */
-    fun updateViewAsync(view: View, attrs: WindowManager.LayoutParams, configuration: Configuration)
+    fun updateViewAsync(
+        view: View,
+        attrs: WindowManager.LayoutParams,
+        configuration: Configuration,
+        touchableRegion: Region? = null,
+    )
 
     /** Releases the underlying [View] hierarchy and removes the backing [SurfaceControl]. */
     fun release(t: SurfaceControl.Transaction)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorWindowlessWindowManager.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorWindowlessWindowManager.kt
new file mode 100644
index 0000000..fbe8c6c8
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/common/viewhost/WindowDecorWindowlessWindowManager.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.windowdecor.common.viewhost
+
+import android.content.res.Configuration
+import android.graphics.Region
+import android.view.SurfaceControl
+import android.view.SurfaceControlViewHost
+import android.view.WindowlessWindowManager
+
+/**
+ * A [WindowlessWindowManager] for the window decor caption that allows customizing the touchable
+ * region.
+ */
+class WindowDecorWindowlessWindowManager(
+    configuration: Configuration,
+    rootSurface: SurfaceControl,
+) : WindowlessWindowManager(configuration, rootSurface, /* hostInputTransferToken= */ null) {
+
+    /** Set the view host's touchable region. */
+    fun setTouchRegion(viewHost: SurfaceControlViewHost, region: Region?) {
+        setTouchRegion(viewHost.windowToken.asBinder(), region)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt
index afd4607..dce44b7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt
@@ -74,6 +74,8 @@
 import com.android.wm.shell.util.StubTransaction
 import com.android.wm.shell.windowdecor.DesktopModeWindowDecorViewModel.DesktopModeKeyguardChangeListener
 import com.android.wm.shell.windowdecor.DesktopModeWindowDecorViewModel.DesktopModeOnInsetsChangedListener
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier
 import com.android.wm.shell.windowdecor.viewholder.AppHeaderViewHolder
 import org.junit.After
 import org.junit.Rule
@@ -131,6 +133,8 @@
     protected val mockAssistContentRequester = mock<AssistContentRequester>()
     protected val bgExecutor = TestShellExecutor()
     protected val mockMultiInstanceHelper = mock<MultiInstanceHelper>()
+    private val mockWindowDecorViewHostSupplier =
+        mock<WindowDecorViewHostSupplier<WindowDecorViewHost>>()
     protected val mockTasksLimiter = mock<DesktopTasksLimiter>()
     protected val mockFreeformTaskTransitionStarter = mock<FreeformTaskTransitionStarter>()
     protected val mockActivityOrientationChangeHandler =
@@ -193,6 +197,7 @@
             mockDesktopImmersiveController,
             mockGenericLinksParser,
             mockAssistContentRequester,
+            mockWindowDecorViewHostSupplier,
             mockMultiInstanceHelper,
             mockDesktopModeWindowDecorFactory,
             mockInputMonitorFactory,
@@ -289,7 +294,7 @@
         whenever(
             mockDesktopModeWindowDecorFactory.create(
                 any(), any(), any(), any(), any(), any(), eq(task), any(), any(), any(), any(),
-                any(), any(), any(), any(), any(), any(), any(), any())
+                any(), any(), any(), any(), any(), any(), any(), any(), any())
         ).thenReturn(decoration)
         decoration.mTaskInfo = task
         whenever(decoration.user).thenReturn(mockUserHandle)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
index 8c77775..5d5d1f2 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
@@ -113,6 +113,8 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.windowdecor.WindowDecoration.RelayoutParams;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 import com.android.wm.shell.windowdecor.viewholder.AppHeaderViewHolder;
 
 import kotlin.Unit;
@@ -185,6 +187,10 @@
     @Mock
     private AttachedSurfaceControl mMockRootSurfaceControl;
     @Mock
+    private WindowDecorViewHostSupplier<WindowDecorViewHost> mMockWindowDecorViewHostSupplier;
+    @Mock
+    private WindowDecorViewHost mMockWindowDecorViewHost;
+    @Mock
     private WindowDecoration.SurfaceControlViewHostFactory mMockSurfaceControlViewHostFactory;
     @Mock
     private TypedArray mMockRoundedCornersRadiusArray;
@@ -274,6 +280,9 @@
                 any())).thenReturn(mMockAppHeaderViewHolder);
         when(mMockDesktopUserRepositories.getCurrent()).thenReturn(mDesktopRepository);
         when(mMockDesktopUserRepositories.getProfile(anyInt())).thenReturn(mDesktopRepository);
+        when(mMockWindowDecorViewHostSupplier.acquire(any(), eq(defaultDisplay)))
+                .thenReturn(mMockWindowDecorViewHost);
+        when(mMockWindowDecorViewHost.getSurfaceControl()).thenReturn(mock(SurfaceControl.class));
     }
 
     @After
@@ -1722,7 +1731,7 @@
                 mMockGenericLinksParser, mMockAssistContentRequester, SurfaceControl.Builder::new,
                 mMockTransactionSupplier, WindowContainerTransaction::new, SurfaceControl::new,
                 new WindowManagerWrapper(mMockWindowManager), mMockSurfaceControlViewHostFactory,
-                maximizeMenuFactory, mMockHandleMenuFactory,
+                mMockWindowDecorViewHostSupplier, maximizeMenuFactory, mMockHandleMenuFactory,
                 mMockMultiInstanceHelper, mMockCaptionHandleRepository, mDesktopModeEventLogger);
         windowDecor.setCaptionListeners(mMockTouchEventListener, mMockTouchEventListener,
                 mMockTouchEventListener, mMockTouchEventListener);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
index 3a82ff6..d969346 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
@@ -89,6 +89,8 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.tests.R;
 import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewContainer;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHost;
+import com.android.wm.shell.windowdecor.common.viewhost.WindowDecorViewHostSupplier;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -131,6 +133,10 @@
     @Mock
     private WindowDecoration.SurfaceControlViewHostFactory mMockSurfaceControlViewHostFactory;
     @Mock
+    private WindowDecorViewHostSupplier<WindowDecorViewHost> mMockWindowDecorViewHostSupplier;
+    @Mock
+    private WindowDecorViewHost mMockWindowDecorViewHost;
+    @Mock
     private SurfaceControlViewHost mMockSurfaceControlViewHost;
     @Mock
     private AttachedSurfaceControl mMockRootSurfaceControl;
@@ -180,6 +186,10 @@
         // Add status bar inset so that WindowDecoration does not think task is in immersive mode
         mInsetsState.getOrCreateSource(STATUS_BAR_INSET_SOURCE_ID, statusBars()).setVisible(true);
         doReturn(mInsetsState).when(mMockDisplayController).getInsetsState(anyInt());
+
+        when(mMockWindowDecorViewHostSupplier.acquire(any(), any()))
+                .thenReturn(mMockWindowDecorViewHost);
+        when(mMockWindowDecorViewHost.getSurfaceControl()).thenReturn(mock(SurfaceControl.class));
     }
 
     @Test
@@ -236,10 +246,6 @@
         final SurfaceControl.Builder decorContainerSurfaceBuilder =
                 createMockSurfaceControlBuilder(decorContainerSurface);
         mMockSurfaceControlBuilders.add(decorContainerSurfaceBuilder);
-        final SurfaceControl captionContainerSurface = mock(SurfaceControl.class);
-        final SurfaceControl.Builder captionContainerSurfaceBuilder =
-                createMockSurfaceControlBuilder(captionContainerSurface);
-        mMockSurfaceControlBuilders.add(captionContainerSurfaceBuilder);
 
         final ActivityManager.RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
                 .setDisplayId(Display.DEFAULT_DISPLAY)
@@ -260,18 +266,19 @@
         verify(mMockSurfaceControlStartT).setTrustedOverlay(decorContainerSurface, true);
         verify(mMockSurfaceControlStartT).setWindowCrop(decorContainerSurface, 300, 100);
 
-        verify(captionContainerSurfaceBuilder).setParent(decorContainerSurface);
-        verify(captionContainerSurfaceBuilder).setContainerLayer();
+        final SurfaceControl captionContainerSurface = mMockWindowDecorViewHost.getSurfaceControl();
+        verify(mMockSurfaceControlStartT).reparent(captionContainerSurface, decorContainerSurface);
         verify(mMockSurfaceControlStartT).setWindowCrop(captionContainerSurface, 300, 64);
         verify(mMockSurfaceControlStartT).show(captionContainerSurface);
 
-        verify(mMockSurfaceControlViewHostFactory).create(any(), eq(defaultDisplay), any());
-
-        verify(mMockSurfaceControlViewHost)
-                .setView(same(mMockView),
-                        argThat(lp -> lp.height == 64
-                                && lp.width == 300
-                                && (lp.flags & LayoutParams.FLAG_NOT_FOCUSABLE) != 0));
+        verify(mMockWindowDecorViewHost).updateView(
+                same(mMockView),
+                argThat(lp -> lp.height == 64
+                        && lp.width == 300
+                        && (lp.flags & LayoutParams.FLAG_NOT_FOCUSABLE) != 0),
+                eq(taskInfo.configuration),
+                any(),
+                eq(null) /* onDrawTransaction */);
         verify(mMockView).setTaskFocusState(true);
         verify(mMockWindowContainerTransaction).addInsetsSource(
                 eq(taskInfo.token),
@@ -300,10 +307,6 @@
         final SurfaceControl.Builder decorContainerSurfaceBuilder =
                 createMockSurfaceControlBuilder(decorContainerSurface);
         mMockSurfaceControlBuilders.add(decorContainerSurfaceBuilder);
-        final SurfaceControl captionContainerSurface = mock(SurfaceControl.class);
-        final SurfaceControl.Builder captionContainerSurfaceBuilder =
-                createMockSurfaceControlBuilder(captionContainerSurface);
-        mMockSurfaceControlBuilders.add(captionContainerSurfaceBuilder);
 
         final SurfaceControl.Transaction t = mock(SurfaceControl.Transaction.class);
         mMockSurfaceControlTransactions.add(t);
@@ -326,7 +329,7 @@
 
         windowDecor.relayout(taskInfo, true /* hasGlobalFocus */);
 
-        verify(mMockSurfaceControlViewHost, never()).release();
+        verify(mMockWindowDecorViewHost, never()).release(any());
         verify(t, never()).apply();
         verify(mMockWindowContainerTransaction, never())
                 .removeInsetsSource(eq(taskInfo.token), any(), anyInt(), anyInt());
@@ -336,9 +339,8 @@
         taskInfo.isVisible = false;
         windowDecor.relayout(taskInfo, false /* hasGlobalFocus */);
 
-        final InOrder releaseOrder = inOrder(t2, mMockSurfaceControlViewHost);
-        releaseOrder.verify(mMockSurfaceControlViewHost).release();
-        releaseOrder.verify(t2).remove(captionContainerSurface);
+        final InOrder releaseOrder = inOrder(t2, mMockWindowDecorViewHostSupplier);
+        releaseOrder.verify(mMockWindowDecorViewHostSupplier).release(mMockWindowDecorViewHost, t2);
         releaseOrder.verify(t2).remove(decorContainerSurface);
         releaseOrder.verify(t2).apply();
         // Expect to remove two insets sources, the caption insets and the mandatory gesture insets.
@@ -386,8 +388,8 @@
         verify(mMockDisplayController).removeDisplayWindowListener(same(listener));
 
         assertThat(mRelayoutResult.mRootView).isSameInstanceAs(mMockView);
-        verify(mMockSurfaceControlViewHostFactory).create(any(), eq(mockDisplay), any());
-        verify(mMockSurfaceControlViewHost).setView(same(mMockView), any());
+        verify(mMockWindowDecorViewHostSupplier).acquire(any(), eq(mockDisplay));
+        verify(mMockWindowDecorViewHost).updateView(same(mMockView), any(), any(), any(), any());
     }
 
     @Test
@@ -400,10 +402,6 @@
         final SurfaceControl.Builder decorContainerSurfaceBuilder =
                 createMockSurfaceControlBuilder(decorContainerSurface);
         mMockSurfaceControlBuilders.add(decorContainerSurfaceBuilder);
-        final SurfaceControl captionContainerSurface = mock(SurfaceControl.class);
-        final SurfaceControl.Builder captionContainerSurfaceBuilder =
-                createMockSurfaceControlBuilder(captionContainerSurface);
-        mMockSurfaceControlBuilders.add(captionContainerSurfaceBuilder);
 
         final SurfaceControl.Transaction t = mock(SurfaceControl.Transaction.class);
         mMockSurfaceControlTransactions.add(t);
@@ -439,8 +437,7 @@
                 windowDecor.mDecorWindowContext.getResources(), mRelayoutParams.mCaptionHeightId);
         verify(mMockSurfaceControlAddWindowT).setWindowCrop(additionalWindowSurface, width, height);
         verify(mMockSurfaceControlAddWindowT).show(additionalWindowSurface);
-        verify(mMockSurfaceControlViewHostFactory, Mockito.times(2))
-                .create(any(), eq(defaultDisplay), any());
+        verify(mMockSurfaceControlViewHostFactory).create(any(), eq(defaultDisplay), any());
     }
 
     @Test
@@ -453,10 +450,6 @@
         final SurfaceControl.Builder decorContainerSurfaceBuilder =
                 createMockSurfaceControlBuilder(decorContainerSurface);
         mMockSurfaceControlBuilders.add(decorContainerSurfaceBuilder);
-        final SurfaceControl captionContainerSurface = mock(SurfaceControl.class);
-        final SurfaceControl.Builder captionContainerSurfaceBuilder =
-                createMockSurfaceControlBuilder(captionContainerSurface);
-        mMockSurfaceControlBuilders.add(captionContainerSurfaceBuilder);
 
         final SurfaceControl.Transaction t = mock(SurfaceControl.Transaction.class);
         mMockSurfaceControlTransactions.add(t);
@@ -475,8 +468,8 @@
 
         windowDecor.relayout(taskInfo, true /* hasGlobalFocus */);
 
-        verify(captionContainerSurfaceBuilder).setParent(decorContainerSurface);
-        verify(captionContainerSurfaceBuilder).setContainerLayer();
+        final SurfaceControl captionContainerSurface = mMockWindowDecorViewHost.getSurfaceControl();
+        verify(mMockSurfaceControlStartT).reparent(captionContainerSurface, decorContainerSurface);
         // Width of the captionContainerSurface should match the width of TASK_BOUNDS
         verify(mMockSurfaceControlStartT).setWindowCrop(captionContainerSurface, 300, 64);
         verify(mMockSurfaceControlStartT).show(captionContainerSurface);
@@ -492,10 +485,6 @@
         final SurfaceControl.Builder decorContainerSurfaceBuilder =
                 createMockSurfaceControlBuilder(decorContainerSurface);
         mMockSurfaceControlBuilders.add(decorContainerSurfaceBuilder);
-        final SurfaceControl captionContainerSurface = mock(SurfaceControl.class);
-        final SurfaceControl.Builder captionContainerSurfaceBuilder =
-                createMockSurfaceControlBuilder(captionContainerSurface);
-        mMockSurfaceControlBuilders.add(captionContainerSurfaceBuilder);
 
         final SurfaceControl.Transaction t = mock(SurfaceControl.Transaction.class);
         mMockSurfaceControlTransactions.add(t);
@@ -512,10 +501,11 @@
         taskInfo.configuration.densityDpi = DisplayMetrics.DENSITY_DEFAULT * 2;
         final TestWindowDecoration windowDecor = createWindowDecoration(taskInfo);
 
-        windowDecor.relayout(taskInfo, true /* applyStartTransactionOnDraw */,
-                true /* hasGlobalFocus */, Region.obtain());
+        mRelayoutParams.mApplyStartTransactionOnDraw = true;
+        windowDecor.relayout(taskInfo, true /* hasGlobalFocus */, Region.obtain());
 
-        verify(mMockRootSurfaceControl).applyTransactionOnDraw(mMockSurfaceControlStartT);
+        verify(mMockWindowDecorViewHost).updateView(any(), any(), any(), any(),
+                eq(mMockSurfaceControlStartT));
     }
 
     @Test
@@ -918,7 +908,13 @@
                 /* hasGlobalFocus= */ true,
                 Region.obtain());
 
-        verify(mMockRootSurfaceControl).applyTransactionOnDraw(mMockSurfaceControlStartT);
+        verify(mMockWindowDecorViewHost)
+                .updateView(
+                        eq(mRelayoutResult.mRootView),
+                        any(),
+                        eq(windowDecor.mTaskInfo.configuration),
+                        any(),
+                        eq(mMockSurfaceControlStartT));
         windowDecor.close();
     }
 
@@ -933,15 +929,11 @@
         mRelayoutParams.mAsyncViewHost = true;
         mRelayoutResult.mRootView = mMockView;
 
-        windowDecor.relayout(
-                windowDecor.mTaskInfo,
-                /* hasGlobalFocus= */ true,
-                Region.obtain());
-        final ArgumentCaptor<Runnable> updateViewHostCaptor =
-                ArgumentCaptor.forClass(Runnable.class);
-        verify(mMockHandler).post(updateViewHostCaptor.capture());
         assertThrows(IllegalArgumentException.class,
-                () -> updateViewHostCaptor.getValue().run());
+                () -> windowDecor.relayout(
+                        windowDecor.mTaskInfo,
+                        /* hasGlobalFocus= */ true,
+                        Region.obtain()));
         windowDecor.close();
     }
 
@@ -961,13 +953,9 @@
                 /* hasGlobalFocus= */ true,
                 Region.obtain());
 
-        final ArgumentCaptor<Runnable> updateViewHostCaptor =
-                ArgumentCaptor.forClass(Runnable.class);
-        verify(mMockHandler).post(updateViewHostCaptor.capture());
-
-        updateViewHostCaptor.getValue().run();
-
-        verify(mMockSurfaceControlViewHost).setView(eq(mMockView), any());
+        verify(mMockWindowDecorViewHost)
+                .updateViewAsync(eq(mRelayoutResult.mRootView), any(),
+                        eq(windowDecor.mTaskInfo.configuration), any());
         windowDecor.close();
     }
 
@@ -1046,13 +1034,14 @@
 
     private TestWindowDecoration createWindowDecoration(ActivityManager.RunningTaskInfo taskInfo) {
         return new TestWindowDecoration(mContext, mContext, mMockDisplayController,
-                mMockShellTaskOrganizer, mMockHandler, taskInfo, mMockTaskSurface,
+                mMockShellTaskOrganizer, taskInfo, mMockTaskSurface,
                 new MockObjectSupplier<>(mMockSurfaceControlBuilders,
                         () -> createMockSurfaceControlBuilder(mock(SurfaceControl.class))),
                 new MockObjectSupplier<>(mMockSurfaceControlTransactions,
                         () -> mock(SurfaceControl.Transaction.class)),
                 () -> mMockWindowContainerTransaction, () -> mMockTaskSurface,
-                mMockSurfaceControlViewHostFactory, mDesktopModeEventLogger);
+                mMockSurfaceControlViewHostFactory, mMockWindowDecorViewHostSupplier,
+                mDesktopModeEventLogger);
     }
 
     private class MockObjectSupplier<T> implements Supplier<T> {
@@ -1087,7 +1076,6 @@
         TestWindowDecoration(Context context, @NonNull Context userContext,
                 DisplayController displayController,
                 ShellTaskOrganizer taskOrganizer,
-                Handler handler,
                 ActivityManager.RunningTaskInfo taskInfo,
                 SurfaceControl taskSurface,
                 Supplier<SurfaceControl.Builder> surfaceControlBuilderSupplier,
@@ -1095,11 +1083,14 @@
                 Supplier<WindowContainerTransaction> windowContainerTransactionSupplier,
                 Supplier<SurfaceControl> surfaceControlSupplier,
                 SurfaceControlViewHostFactory surfaceControlViewHostFactory,
+                @NonNull WindowDecorViewHostSupplier<WindowDecorViewHost>
+                        windowDecorViewHostSupplier,
                 DesktopModeEventLogger desktopModeEventLogger) {
-            super(context, userContext, displayController, taskOrganizer, handler, taskInfo,
+            super(context, userContext, displayController, taskOrganizer, taskInfo,
                     taskSurface, surfaceControlBuilderSupplier, surfaceControlTransactionSupplier,
                     windowContainerTransactionSupplier, surfaceControlSupplier,
-                    surfaceControlViewHostFactory, desktopModeEventLogger);
+                    surfaceControlViewHostFactory, windowDecorViewHostSupplier,
+                    desktopModeEventLogger);
         }
 
         void relayout(ActivityManager.RunningTaskInfo taskInfo, boolean hasGlobalFocus) {
diff --git a/media/java/android/media/MediaRoute2Info.java b/media/java/android/media/MediaRoute2Info.java
index dd5067a..54f172c 100644
--- a/media/java/android/media/MediaRoute2Info.java
+++ b/media/java/android/media/MediaRoute2Info.java
@@ -189,6 +189,7 @@
      * the device.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_BUILTIN_SPEAKER
      */
     public static final int TYPE_BUILTIN_SPEAKER = AudioDeviceInfo.TYPE_BUILTIN_SPEAKER;
 
@@ -196,6 +197,7 @@
      * Indicates the route is a headset, which is the combination of a headphones and a microphone.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_WIRED_HEADSET
      */
     public static final int TYPE_WIRED_HEADSET = AudioDeviceInfo.TYPE_WIRED_HEADSET;
 
@@ -203,6 +205,7 @@
      * Indicates the route is a pair of wired headphones.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_WIRED_HEADPHONES
      */
     public static final int TYPE_WIRED_HEADPHONES = AudioDeviceInfo.TYPE_WIRED_HEADPHONES;
 
@@ -210,6 +213,7 @@
      * Indicates the route is a bluetooth device, such as a bluetooth speaker or headphones.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_BLUETOOTH_A2DP
      */
     public static final int TYPE_BLUETOOTH_A2DP = AudioDeviceInfo.TYPE_BLUETOOTH_A2DP;
 
@@ -217,6 +221,7 @@
      * Indicates the route is an HDMI connection.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_HDMI
      */
     public static final int TYPE_HDMI = AudioDeviceInfo.TYPE_HDMI;
 
@@ -224,6 +229,7 @@
      * Indicates the route is an Audio Return Channel of an HDMI connection.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_HDMI_ARC
      */
     @FlaggedApi(FLAG_ENABLE_AUDIO_POLICIES_DEVICE_AND_BLUETOOTH_CONTROLLER)
     public static final int TYPE_HDMI_ARC = AudioDeviceInfo.TYPE_HDMI_ARC;
@@ -232,24 +238,34 @@
      * Indicates the route is an Enhanced Audio Return Channel of an HDMI connection.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_HDMI_EARC
      */
     @FlaggedApi(FLAG_ENABLE_AUDIO_POLICIES_DEVICE_AND_BLUETOOTH_CONTROLLER)
     public static final int TYPE_HDMI_EARC = AudioDeviceInfo.TYPE_HDMI_EARC;
 
     /**
      * Indicates the route is a digital line connection (for example S/PDIF).
+     *
+     * @see #getType
+     * @see AudioDeviceInfo#TYPE_LINE_DIGITAL
      */
     @FlaggedApi(FLAG_ENABLE_NEW_WIRED_MEDIA_ROUTE_2_INFO_TYPES)
     public static final int TYPE_LINE_DIGITAL = AudioDeviceInfo.TYPE_LINE_DIGITAL;
 
     /**
      * Indicates the route is an analog line-level connection.
+     *
+     * @see #getType
+     * @see AudioDeviceInfo#TYPE_LINE_ANALOG
      */
     @FlaggedApi(FLAG_ENABLE_NEW_WIRED_MEDIA_ROUTE_2_INFO_TYPES)
     public static final int TYPE_LINE_ANALOG = AudioDeviceInfo.TYPE_LINE_ANALOG;
 
     /**
      * Indicates the route is using the auxiliary line-level connectors.
+     *
+     * @see #getType
+     * @see AudioDeviceInfo#TYPE_AUX_LINE
      */
     @FlaggedApi(FLAG_ENABLE_NEW_WIRED_MEDIA_ROUTE_2_INFO_TYPES)
     public static final int TYPE_AUX_LINE = AudioDeviceInfo.TYPE_AUX_LINE;
@@ -258,6 +274,7 @@
      * Indicates the route is a USB audio device.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_USB_DEVICE
      */
     public static final int TYPE_USB_DEVICE = AudioDeviceInfo.TYPE_USB_DEVICE;
 
@@ -265,6 +282,7 @@
      * Indicates the route is a USB audio device in accessory mode.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_USB_ACCESSORY
      */
     public static final int TYPE_USB_ACCESSORY = AudioDeviceInfo.TYPE_USB_ACCESSORY;
 
@@ -272,6 +290,7 @@
      * Indicates the route is the audio device associated with a dock.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_DOCK
      */
     public static final int TYPE_DOCK = AudioDeviceInfo.TYPE_DOCK;
 
@@ -279,6 +298,7 @@
      * Indicates the route is a USB audio headset.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_USB_HEADSET
      */
     public static final int TYPE_USB_HEADSET = AudioDeviceInfo.TYPE_USB_HEADSET;
 
@@ -286,6 +306,7 @@
      * Indicates the route is a hearing aid.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_HEARING_AID
      */
     public static final int TYPE_HEARING_AID = AudioDeviceInfo.TYPE_HEARING_AID;
 
@@ -293,6 +314,7 @@
      * Indicates the route is a Bluetooth Low Energy (BLE) HEADSET.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_BLE_HEADSET
      */
     public static final int TYPE_BLE_HEADSET = AudioDeviceInfo.TYPE_BLE_HEADSET;
 
@@ -304,6 +326,7 @@
      * to provide a better experience on multichannel contents.
      *
      * @see #getType
+     * @see AudioDeviceInfo#TYPE_MULTICHANNEL_GROUP
      */
     @FlaggedApi(FLAG_ENABLE_MULTICHANNEL_GROUP_DEVICE)
     public static final int TYPE_MULTICHANNEL_SPEAKER_GROUP =
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt
index cd03dd7..07b1c9e 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt
@@ -22,6 +22,7 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue
 import java.util.WeakHashMap
 import java.util.concurrent.Executor
+import java.util.concurrent.atomic.AtomicInteger
 
 /**
  * Callback to be informed of changes in [KeyedObservable] object.
@@ -203,13 +204,71 @@
         }
     }
 
-    fun hasAnyObserver(): Boolean {
+    open fun hasAnyObserver(): Boolean {
         synchronized(observers) { if (observers.isNotEmpty()) return true }
         synchronized(keyedObservers) { if (keyedObservers.isNotEmpty()) return true }
         return false
     }
 }
 
+/** [KeyedDataObservable] that maintains a counter for the observers. */
+abstract class AbstractKeyedDataObservable<K> : KeyedDataObservable<K>() {
+    /**
+     * Counter of observers.
+     *
+     * The value is accurate only when [addObserver] and [removeObserver] are invoked in pairs.
+     */
+    private val counter = AtomicInteger()
+
+    override fun addObserver(observer: KeyedObserver<K?>, executor: Executor) =
+        if (super.addObserver(observer, executor)) {
+            onObserverAdded()
+            true
+        } else {
+            false
+        }
+
+    override fun addObserver(key: K, observer: KeyedObserver<K>, executor: Executor) =
+        if (super.addObserver(key, observer, executor)) {
+            onObserverAdded()
+            true
+        } else {
+            false
+        }
+
+    private fun onObserverAdded() {
+        if (counter.getAndIncrement() == 0) onFirstObserverAdded()
+    }
+
+    /** Callbacks when the first observer is just added. */
+    protected abstract fun onFirstObserverAdded()
+
+    override fun removeObserver(observer: KeyedObserver<K?>) =
+        if (super.removeObserver(observer)) {
+            onObserverRemoved()
+            true
+        } else {
+            false
+        }
+
+    override fun removeObserver(key: K, observer: KeyedObserver<K>) =
+        if (super.removeObserver(key, observer)) {
+            onObserverRemoved()
+            true
+        } else {
+            false
+        }
+
+    private fun onObserverRemoved() {
+        if (counter.decrementAndGet() == 0) onLastObserverRemoved()
+    }
+
+    /** Callbacks when the last observer is just removed. */
+    protected abstract fun onLastObserverRemoved()
+
+    override fun hasAnyObserver() = counter.get() > 0
+}
+
 /** [KeyedObservable] with no-op implementations for all interfaces. */
 open class NoOpKeyedObservable<K> : KeyedObservable<K> {
 
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SettingsStore.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SettingsStore.kt
index 04d4bfe..d6e7a89 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SettingsStore.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SettingsStore.kt
@@ -20,21 +20,10 @@
 import android.database.ContentObserver
 import android.net.Uri
 import android.util.Log
-import java.util.concurrent.Executor
-import java.util.concurrent.atomic.AtomicInteger
 
 /** Base class of the Settings provider data stores. */
 abstract class SettingsStore(protected val contentResolver: ContentResolver) :
-    KeyedDataObservable<String>(), KeyValueStore {
-
-    /**
-     * Counter of observers.
-     *
-     * The value is accurate only when [addObserver] and [removeObserver] are called correctly. When
-     * an observer is not removed (and its weak reference is garbage collected), the content
-     * observer is not unregistered but this is not a big deal.
-     */
-    private val counter = AtomicInteger()
+    AbstractKeyedDataObservable<String>(), KeyValueStore {
 
     private val contentObserver =
         object : ContentObserver(HandlerExecutor.main) {
@@ -48,49 +37,15 @@
             }
         }
 
-    override fun addObserver(observer: KeyedObserver<String?>, executor: Executor) =
-        if (super.addObserver(observer, executor)) {
-            onObserverAdded()
-            true
-        } else {
-            false
-        }
+    /** The URI to watch for any key change. */
+    protected abstract val uri: Uri
 
-    override fun addObserver(key: String, observer: KeyedObserver<String>, executor: Executor) =
-        if (super.addObserver(key, observer, executor)) {
-            onObserverAdded()
-            true
-        } else {
-            false
-        }
-
-    private fun onObserverAdded() {
-        if (counter.getAndIncrement() != 0) return
+    override fun onFirstObserverAdded() {
         Log.i(tag, "registerContentObserver")
         contentResolver.registerContentObserver(uri, true, contentObserver)
     }
 
-    /** The URI to watch for any key change. */
-    protected abstract val uri: Uri
-
-    override fun removeObserver(observer: KeyedObserver<String?>) =
-        if (super.removeObserver(observer)) {
-            onObserverRemoved()
-            true
-        } else {
-            false
-        }
-
-    override fun removeObserver(key: String, observer: KeyedObserver<String>) =
-        if (super.removeObserver(key, observer)) {
-            onObserverRemoved()
-            true
-        } else {
-            false
-        }
-
-    private fun onObserverRemoved() {
-        if (counter.decrementAndGet() != 0) return
+    override fun onLastObserverRemoved() {
         Log.i(tag, "unregisterContentObserver")
         contentResolver.unregisterContentObserver(contentObserver)
     }