Fix CUJ UiThread detection when only a SurfaceControl is provided

Most methods annotated with @UiThread were not actually running on the UI thread related to the CUJ (when its ui thread was different than the app main thread).

When deferred monitoring was true (the default case), everything was fine, as the "begin" method eventually goes to the choreographer, and is scheduled on the correct UI thread.

When deferred monitoring was true and only a surface was provided though, the "begin" method was being executed

Bug: 363884280
Test: FrameTrackerTest and Checked ui thread associated with splashscreen CUJs and fold animation
Flag: NONE recent bugfix
Change-Id: I3aa1e67b50d47f1eb6ae5c896e53f6955298e750

34 files changed, 317 insertions, 103 deletions

diff --git a/core/java/android/view/inputmethod/ImeTracker.java b/core/java/android/view/inputmethod/ImeTracker.java
index b9751c8cab52..d90455ab971d 100644
--- a/core/java/android/view/inputmethod/ImeTracker.java
+++ b/core/java/android/view/inputmethod/ImeTracker.java
@@ -920,7 +920,8 @@ public interface ImeTracker {
             final Configuration.Builder builder = Configuration.Builder.withSurface(
                             cujType,
                             jankContext.getDisplayContext(),
-                            jankContext.getTargetSurfaceControl())
+                            jankContext.getTargetSurfaceControl(),
+                            jankContext.getDisplayContext().getMainThreadHandler())
                     .setTag(String.format(Locale.US, "%d@%d@%s", animType,
                             useSeparatedThread ? 0 : 1, jankContext.getHostPackageName()));
             InteractionJankMonitor.getInstance().begin(builder);
diff --git a/core/java/com/android/internal/jank/InteractionJankMonitor.java b/core/java/com/android/internal/jank/InteractionJankMonitor.java
index c7e1fba66d7f..ef08e49ce6d9 100644
--- a/core/java/com/android/internal/jank/InteractionJankMonitor.java
+++ b/core/java/com/android/internal/jank/InteractionJankMonitor.java
@@ -330,9 +330,10 @@ public class InteractionJankMonitor {
      * @param cujType the specific {@link Cuj.CujType}.
      * @return boolean true if the tracker is started successfully, false otherwise.
      */
-    public boolean begin(SurfaceControl surface, Context context, @Cuj.CujType int cujType) {
+    public boolean begin(SurfaceControl surface, Context context, Handler handler,
+            @Cuj.CujType int cujType) {
         try {
-            return begin(Configuration.Builder.withSurface(cujType, context, surface));
+            return begin(Configuration.Builder.withSurface(cujType, context, surface, handler));
         } catch (IllegalArgumentException ex) {
             Log.d(TAG, "Build configuration failed!", ex);
             return false;
@@ -348,11 +349,12 @@ public class InteractionJankMonitor {
      * @param tag a tag containing extra information about the interaction.
      * @return boolean true if the tracker is started successfully, false otherwise.
      */
-    public boolean begin(SurfaceControl surface, Context context, @Cuj.CujType int cujType,
+    public boolean begin(SurfaceControl surface, Context context, Handler handler,
+            @Cuj.CujType int cujType,
             String tag) {
         try {
             final Configuration.Builder builder =
-                    Configuration.Builder.withSurface(cujType, context, surface);
+                    Configuration.Builder.withSurface(cujType, context, surface, handler);
             if (!TextUtils.isEmpty(tag)) {
                 builder.setTag(tag);
             }
@@ -689,20 +691,23 @@ public class InteractionJankMonitor {
             private SurfaceControl mAttrSurfaceControl;
             private final @Cuj.CujType int mAttrCujType;
             private boolean mAttrDeferMonitor = true;
+            private Handler mHandler = null;
 
             /**
              * Creates a builder which instruments only surface.
              * @param cuj The enum defined in {@link Cuj.CujType}.
              * @param context context
              * @param surfaceControl surface control
+             * @param uiThreadHandler UI thread for that surface
              * @return builder
              */
             public static Builder withSurface(@Cuj.CujType int cuj, @NonNull Context context,
-                    @NonNull SurfaceControl surfaceControl) {
+                    @NonNull SurfaceControl surfaceControl, @NonNull Handler uiThreadHandler) {
                 return new Builder(cuj)
                         .setContext(context)
                         .setSurfaceControl(surfaceControl)
-                        .setSurfaceOnly(true);
+                        .setSurfaceOnly(true)
+                        .setHandler(uiThreadHandler);
             }
 
             /**
@@ -722,6 +727,18 @@ public class InteractionJankMonitor {
             }
 
             /**
+             * Specifies the UI thread handler. If not provided, the View's one will be used.
+             * If only a surface is provided without handler, the app main thread will be used.
+             *
+             * @param uiThreadHandler handler associated to the cuj UI thread
+             * @return builder
+             */
+            public Builder setHandler(Handler uiThreadHandler) {
+                mHandler = uiThreadHandler;
+                return this;
+            }
+
+            /**
              * Specifies a view, must be set if {@link #setSurfaceOnly(boolean)} is set to false.
              * @param view an attached view
              * @return builder
@@ -798,13 +815,13 @@ public class InteractionJankMonitor {
                 return new Configuration(
                         mAttrCujType, mAttrView, mAttrTag, mAttrTimeout,
                         mAttrSurfaceOnly, mAttrContext, mAttrSurfaceControl,
-                        mAttrDeferMonitor);
+                        mAttrDeferMonitor, mHandler);
             }
         }
 
         private Configuration(@Cuj.CujType int cuj, View view, @NonNull String tag, long timeout,
                 boolean surfaceOnly, Context context, SurfaceControl surfaceControl,
-                boolean deferMonitor) {
+                boolean deferMonitor, Handler handler) {
             mCujType = cuj;
             mTag = tag;
             mSessionName = generateSessionName(Cuj.getNameOfCuj(cuj), tag);
@@ -816,8 +833,16 @@ public class InteractionJankMonitor {
                     : (view != null ? view.getContext().getApplicationContext() : null);
             mSurfaceControl = surfaceControl;
             mDeferMonitor = deferMonitor;
+            if (handler != null) {
+                mHandler = handler;
+            } else if (mSurfaceOnly) {
+                Log.w(TAG, "No UIThread provided for " + mSessionName
+                        + " (surface only). Defaulting to app main thread.");
+                mHandler = mContext.getMainThreadHandler();
+            } else {
+                mHandler = mView.getHandler();
+            }
             validate();
-            mHandler = mSurfaceOnly ? mContext.getMainThreadHandler() : mView.getHandler();
         }
 
         @VisibleForTesting
@@ -858,6 +883,12 @@ public class InteractionJankMonitor {
                     shouldThrow = true;
                     msg.append("Must pass in a valid surface control if only instrument surface; ");
                 }
+                if (mHandler == null) {
+                    shouldThrow = true;
+                    msg.append(
+                            "Must pass a UI thread handler when only a surface control is "
+                                    + "provided.");
+                }
             } else {
                 if (!hasValidView()) {
                     shouldThrow = true;
diff --git a/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java b/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java
index 5af272c8bead..e0823b874edf 100644
--- a/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java
+++ b/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java
@@ -24,6 +24,7 @@ import static com.android.internal.jank.InteractionJankMonitor.Configuration.gen
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -117,6 +118,7 @@ public class InteractionJankMonitorTest {
 
         // Simulate a trace session and see if begin / end are invoked.
         assertThat(monitor.begin(mSurfaceControl, mActivity.getApplicationContext(),
+                mActivity.getMainThreadHandler(),
                 Cuj.CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
         verify(tracker).begin();
         assertThat(monitor.end(Cuj.CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
@@ -188,6 +190,25 @@ public class InteractionJankMonitorTest {
         assertThat(generateSessionName(cujName, tooLongTag)).isEqualTo(expectedTrimmedName);
     }
 
+    @Test
+    public void validateConfiguration_surfaceOnlyAndNotDeferMonitor_throwsError() {
+        Configuration.Builder builder = Configuration.Builder.withSurface(1,
+                mActivity.getApplicationContext(), mSurfaceControl,
+                mActivity.getMainThreadHandler()).setDeferMonitorForAnimationStart(false);
+
+        assertThrows(IllegalStateException.class, builder::build);
+    }
+
+    @Test
+    public void validateConfiguration_surfaceOnlyAndDeferMonitor_doesNotThrowError() {
+        Configuration.Builder builder = Configuration.Builder.withSurface(1,
+                mActivity.getApplicationContext(),
+                mSurfaceControl, mActivity.getMainThreadHandler()).setDeferMonitorForAnimationStart(
+                true);
+
+        builder.build(); // no exception.
+    }
+
     private InteractionJankMonitor createMockedInteractionJankMonitor() {
         InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
         doReturn(true).when(monitor).shouldMonitor();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
index f478b4446cbe..3e5adf395cdd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
@@ -117,6 +117,7 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
      */
     private static final long MAX_ANIMATION_DURATION = 2000;
     private final LatencyTracker mLatencyTracker;
+    @ShellMainThread private final Handler mHandler;
 
     /** True when a back gesture is ongoing */
     private boolean mBackGestureStarted = false;
@@ -218,7 +219,8 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
             @NonNull BackAnimationBackground backAnimationBackground,
             ShellBackAnimationRegistry shellBackAnimationRegistry,
             ShellCommandHandler shellCommandHandler,
-            Transitions transitions) {
+            Transitions transitions,
+            @ShellMainThread Handler handler) {
         this(
                 shellInit,
                 shellController,
@@ -230,7 +232,8 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
                 backAnimationBackground,
                 shellBackAnimationRegistry,
                 shellCommandHandler,
-                transitions);
+                transitions,
+                handler);
     }
 
     @VisibleForTesting
@@ -245,7 +248,8 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
             @NonNull BackAnimationBackground backAnimationBackground,
             ShellBackAnimationRegistry shellBackAnimationRegistry,
             ShellCommandHandler shellCommandHandler,
-            Transitions transitions) {
+            Transitions transitions,
+            @NonNull @ShellMainThread Handler handler) {
         mShellController = shellController;
         mShellExecutor = shellExecutor;
         mActivityTaskManager = activityTaskManager;
@@ -263,6 +267,7 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
         mTransitions = transitions;
         mBackTransitionHandler = new BackTransitionHandler();
         mTransitions.addHandler(mBackTransitionHandler);
+        mHandler = handler;
         updateTouchableArea();
     }
 
@@ -399,7 +404,7 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
         }
     }
 
-    private static class IBackAnimationImpl extends IBackAnimation.Stub
+    private class IBackAnimationImpl extends IBackAnimation.Stub
             implements ExternalInterfaceBinder {
         private BackAnimationController mController;
 
@@ -417,7 +422,8 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
                                     callback,
                                     runner,
                                     controller.mContext,
-                                    CUJ_PREDICTIVE_BACK_HOME)));
+                                    CUJ_PREDICTIVE_BACK_HOME,
+                                    mHandler)));
         }
 
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationRunner.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationRunner.java
index e24df0bdc05d..9ca9b730fb06 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationRunner.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationRunner.java
@@ -20,6 +20,7 @@ import static android.view.WindowManager.TRANSIT_OLD_UNSET;
 
 import android.annotation.NonNull;
 import android.content.Context;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.util.Log;
 import android.view.IRemoteAnimationFinishedCallback;
@@ -31,6 +32,7 @@ import android.window.IOnBackInvokedCallback;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.jank.Cuj.CujType;
 import com.android.internal.jank.InteractionJankMonitor;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 
 /**
  * Used to register the animation callback and runner, it will trigger result if gesture was finish
@@ -45,6 +47,8 @@ public class BackAnimationRunner {
     private final IRemoteAnimationRunner mRunner;
     private final @CujType int mCujType;
     private final Context mContext;
+    @ShellMainThread
+    private final Handler mHandler;
 
     // Whether we are waiting to receive onAnimationStart
     private boolean mWaitingAnimation;
@@ -56,18 +60,35 @@ public class BackAnimationRunner {
             @NonNull IOnBackInvokedCallback callback,
             @NonNull IRemoteAnimationRunner runner,
             @NonNull Context context,
-            @CujType int cujType) {
+            @CujType int cujType,
+            @ShellMainThread Handler handler) {
         mCallback = callback;
         mRunner = runner;
         mCujType = cujType;
         mContext = context;
+        mHandler = handler;
     }
 
     public BackAnimationRunner(
             @NonNull IOnBackInvokedCallback callback,
             @NonNull IRemoteAnimationRunner runner,
-            @NonNull Context context) {
-        this(callback, runner, context, NO_CUJ);
+            @NonNull Context context,
+            @ShellMainThread Handler handler
+    ) {
+        this(callback, runner, context, NO_CUJ, handler);
+    }
+
+    /**
+     * @deprecated Use {@link BackAnimationRunner} constructor providing an handler for the ui
+     * thread of the animation.
+     */
+    @Deprecated
+    public BackAnimationRunner(
+            @NonNull IOnBackInvokedCallback callback,
+            @NonNull IRemoteAnimationRunner runner,
+            @NonNull Context context
+    ) {
+        this(callback, runner, context, NO_CUJ, context.getMainThreadHandler());
     }
 
     /** Returns the registered animation runner */
@@ -100,7 +121,7 @@ public class BackAnimationRunner {
         mWaitingAnimation = false;
         if (shouldMonitorCUJ(apps)) {
             interactionJankMonitor.begin(
-                    apps[0].leash, mContext, mCujType);
+                    apps[0].leash, mContext, mHandler, mCujType);
         }
         try {
             getRunner().onAnimationStart(TRANSIT_OLD_UNSET, apps, wallpapers,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt
index 32e809a23a91..37339307f5b0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt
@@ -26,6 +26,7 @@ import android.graphics.Matrix
 import android.graphics.PointF
 import android.graphics.Rect
 import android.graphics.RectF
+import android.os.Handler
 import android.os.RemoteException
 import android.util.TimeUtils
 import android.view.Choreographer
@@ -53,6 +54,7 @@ import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.protolog.ShellProtoLogGroup
 import com.android.wm.shell.shared.animation.Interpolators
+import com.android.wm.shell.shared.annotations.ShellMainThread
 import kotlin.math.abs
 import kotlin.math.max
 import kotlin.math.min
@@ -61,7 +63,8 @@ abstract class CrossActivityBackAnimation(
     private val context: Context,
     private val background: BackAnimationBackground,
     private val rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer,
-    protected val transaction: SurfaceControl.Transaction
+    protected val transaction: SurfaceControl.Transaction,
+    @ShellMainThread handler: Handler,
 ) : ShellBackAnimation() {
 
     protected val startClosingRect = RectF()
@@ -80,7 +83,13 @@ abstract class CrossActivityBackAnimation(
     private var statusbarHeight = SystemBarUtils.getStatusBarHeight(context)
 
     private val backAnimationRunner =
-        BackAnimationRunner(Callback(), Runner(), context, Cuj.CUJ_PREDICTIVE_BACK_CROSS_ACTIVITY)
+        BackAnimationRunner(
+            Callback(),
+            Runner(),
+            context,
+            Cuj.CUJ_PREDICTIVE_BACK_CROSS_ACTIVITY,
+            handler,
+        )
     private val initialTouchPos = PointF()
     private val transformMatrix = Matrix()
     private val tmpFloat9 = FloatArray(9)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java
index 3fccecab4319..7a569799ab84 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java
@@ -34,6 +34,7 @@ import android.graphics.Matrix;
 import android.graphics.PointF;
 import android.graphics.Rect;
 import android.graphics.RectF;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.view.Choreographer;
 import android.view.IRemoteAnimationFinishedCallback;
@@ -52,6 +53,7 @@ import com.android.internal.policy.SystemBarUtils;
 import com.android.internal.protolog.ProtoLog;
 import com.android.wm.shell.R;
 import com.android.wm.shell.shared.animation.Interpolators;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 
 import javax.inject.Inject;
 
@@ -113,9 +115,10 @@ public class CrossTaskBackAnimation extends ShellBackAnimation {
     private float mVerticalMargin;
 
     @Inject
-    public CrossTaskBackAnimation(Context context, BackAnimationBackground background) {
+    public CrossTaskBackAnimation(Context context, BackAnimationBackground background,
+            @ShellMainThread Handler handler) {
         mBackAnimationRunner = new BackAnimationRunner(
-                new Callback(), new Runner(), context, CUJ_PREDICTIVE_BACK_CROSS_TASK);
+                new Callback(), new Runner(), context, CUJ_PREDICTIVE_BACK_CROSS_TASK, handler);
         mBackground = background;
         mContext = context;
         loadResources();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CustomCrossActivityBackAnimation.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CustomCrossActivityBackAnimation.kt
index b02f97bf7784..2f7666b21882 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CustomCrossActivityBackAnimation.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CustomCrossActivityBackAnimation.kt
@@ -18,6 +18,7 @@ package com.android.wm.shell.back
 import android.content.Context
 import android.graphics.Rect
 import android.graphics.RectF
+import android.os.Handler
 import android.util.MathUtils
 import android.view.SurfaceControl
 import android.view.animation.Animation
@@ -30,6 +31,7 @@ import com.android.internal.policy.TransitionAnimation
 import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.protolog.ShellProtoLogGroup
+import com.android.wm.shell.shared.annotations.ShellMainThread
 import javax.inject.Inject
 import kotlin.math.max
 import kotlin.math.min
@@ -40,13 +42,15 @@ class CustomCrossActivityBackAnimation(
     background: BackAnimationBackground,
     rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer,
     transaction: SurfaceControl.Transaction,
-    private val customAnimationLoader: CustomAnimationLoader
+    private val customAnimationLoader: CustomAnimationLoader,
+    @ShellMainThread handler: Handler,
 ) :
     CrossActivityBackAnimation(
         context,
         background,
         rootTaskDisplayAreaOrganizer,
-        transaction
+        transaction,
+        handler
     ) {
 
     private var enterAnimation: Animation? = null
@@ -59,7 +63,8 @@ class CustomCrossActivityBackAnimation(
     constructor(
         context: Context,
         background: BackAnimationBackground,
-        rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
+        rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer,
+        @ShellMainThread handler: Handler,
     ) : this(
         context,
         background,
@@ -67,7 +72,8 @@ class CustomCrossActivityBackAnimation(
         SurfaceControl.Transaction(),
         CustomAnimationLoader(
             TransitionAnimation(context, false /* debug */, "CustomCrossActivityBackAnimation")
-        )
+        ),
+        handler,
     )
 
     override fun preparePreCommitClosingRectMovement(swipeEdge: Int) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/DefaultCrossActivityBackAnimation.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/back/DefaultCrossActivityBackAnimation.kt
index 66d8a5f2eeb9..eecd7694009d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/DefaultCrossActivityBackAnimation.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/DefaultCrossActivityBackAnimation.kt
@@ -16,11 +16,13 @@
 package com.android.wm.shell.back
 
 import android.content.Context
+import android.os.Handler
 import android.view.SurfaceControl
 import android.window.BackEvent
 import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.shared.animation.Interpolators
+import com.android.wm.shell.shared.annotations.ShellMainThread
 import javax.inject.Inject
 import kotlin.math.max
 
@@ -30,13 +32,15 @@ class DefaultCrossActivityBackAnimation
 constructor(
     context: Context,
     background: BackAnimationBackground,
-    rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
+    rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer,
+    @ShellMainThread handler: Handler,
 ) :
     CrossActivityBackAnimation(
         context,
         background,
         rootTaskDisplayAreaOrganizer,
-        SurfaceControl.Transaction()
+        SurfaceControl.Transaction(),
+        handler
     ) {
 
     private val postCommitInterpolator = Interpolators.EMPHASIZED
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
index b8aa1b189f7e..4b55fd0f5527 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
@@ -48,6 +48,7 @@ import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.view.Display;
 import android.view.InsetsSourceControl;
 import android.view.InsetsState;
@@ -75,6 +76,7 @@ import com.android.wm.shell.common.pip.PipUtils;
 import com.android.wm.shell.common.split.DividerSnapAlgorithm.SnapTarget;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.shared.animation.Interpolators;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.shared.split.SplitScreenConstants.PersistentSnapPosition;
 import com.android.wm.shell.shared.split.SplitScreenConstants.SnapPosition;
 import com.android.wm.shell.shared.split.SplitScreenConstants.SplitPosition;
@@ -116,6 +118,8 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
             new PathInterpolator(0.2f, 0f, 0f, 1f);
     private static final Interpolator GROW_INTERPOLATOR =
             new PathInterpolator(0.45f, 0f, 0.5f, 1f);
+    @ShellMainThread
+    private final Handler mHandler;
 
     private int mDividerWindowWidth;
     private int mDividerInsets;
@@ -166,7 +170,8 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
             SplitLayoutHandler splitLayoutHandler,
             SplitWindowManager.ParentContainerCallbacks parentContainerCallbacks,
             DisplayController displayController, DisplayImeController displayImeController,
-            ShellTaskOrganizer taskOrganizer, int parallaxType) {
+            ShellTaskOrganizer taskOrganizer, int parallaxType, @ShellMainThread Handler handler) {
+        mHandler = handler;
         mContext = context.createConfigurationContext(configuration);
         mOrientation = configuration.orientation;
         mRotation = configuration.windowConfiguration.getRotation();
@@ -598,7 +603,8 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
     }
 
     void onStartDragging() {
-        mInteractionJankMonitor.begin(getDividerLeash(), mContext, CUJ_SPLIT_SCREEN_RESIZE);
+        mInteractionJankMonitor.begin(getDividerLeash(), mContext, mHandler,
+                CUJ_SPLIT_SCREEN_RESIZE);
     }
 
     void onDraggingCancelled() {
@@ -756,7 +762,7 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
             @Override
             public void onAnimationStart(Animator animation) {
                 mInteractionJankMonitor.begin(getDividerLeash(),
-                        mContext, CUJ_SPLIT_SCREEN_DOUBLE_TAP_DIVIDER);
+                        mContext, mHandler, CUJ_SPLIT_SCREEN_DOUBLE_TAP_DIVIDER);
             }
 
             @Override
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 4adea233b734..722fe1f59388 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
@@ -444,7 +444,9 @@ public abstract class WMShellBaseModule {
             BackAnimationBackground backAnimationBackground,
             Optional<ShellBackAnimationRegistry> shellBackAnimationRegistry,
             ShellCommandHandler shellCommandHandler,
-            Transitions transitions) {
+            Transitions transitions,
+            @ShellMainThread Handler handler
+    ) {
         if (BackAnimationController.IS_ENABLED) {
             return shellBackAnimationRegistry.map(
                     (animations) ->
@@ -457,7 +459,8 @@ public abstract class WMShellBaseModule {
                                     backAnimationBackground,
                                     animations,
                                     shellCommandHandler,
-                                    transitions));
+                                    transitions,
+                                    handler));
         }
         return Optional.empty();
     }
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 b151c8b7e718..cf1d99e06d9b 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
@@ -355,7 +355,8 @@ public abstract class WMShellModule {
             @ShellMainThread ShellExecutor mainExecutor,
             @ShellAnimationThread ShellExecutor animExecutor,
             @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
-            InteractionJankMonitor interactionJankMonitor) {
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler) {
         return new FreeformTaskTransitionHandler(
                 shellInit,
                 transitions,
@@ -365,7 +366,8 @@ public abstract class WMShellModule {
                 mainExecutor,
                 animExecutor,
                 desktopModeTaskRepository,
-                interactionJankMonitor);
+                interactionJankMonitor,
+                handler);
     }
 
     @WMSingleton
@@ -616,6 +618,7 @@ public abstract class WMShellModule {
             RecentsTransitionHandler recentsTransitionHandler,
             MultiInstanceHelper multiInstanceHelper,
             @ShellMainThread ShellExecutor mainExecutor,
+            @ShellMainThread Handler mainHandler,
             Optional<DesktopTasksLimiter> desktopTasksLimiter,
             Optional<RecentTasksController> recentTasksController,
             InteractionJankMonitor interactionJankMonitor) {
@@ -628,7 +631,7 @@ public abstract class WMShellModule {
                 dragToDesktopTransitionHandler, desktopModeTaskRepository,
                 desktopModeLoggerTransitionObserver, launchAdjacentController,
                 recentsTransitionHandler, multiInstanceHelper, mainExecutor, desktopTasksLimiter,
-                recentTasksController.orElse(null), interactionJankMonitor);
+                recentTasksController.orElse(null), interactionJankMonitor, mainHandler);
     }
 
     @WMSingleton
@@ -638,7 +641,8 @@ public abstract class WMShellModule {
             Transitions transitions,
             @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
             ShellTaskOrganizer shellTaskOrganizer,
-            InteractionJankMonitor interactionJankMonitor) {
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler) {
         int maxTaskLimit = DesktopModeStatus.getMaxTaskLimit(context);
         if (!DesktopModeStatus.canEnterDesktopMode(context)
                 || !ENABLE_DESKTOP_WINDOWING_TASK_LIMIT.isEnabled(context)
@@ -652,7 +656,8 @@ public abstract class WMShellModule {
                         shellTaskOrganizer,
                         maxTaskLimit,
                         interactionJankMonitor,
-                        context)
+                        context,
+                        handler)
         );
     }
 
@@ -699,9 +704,10 @@ public abstract class WMShellModule {
     static ExitDesktopTaskTransitionHandler provideExitDesktopTaskTransitionHandler(
             Transitions transitions,
             Context context,
-            InteractionJankMonitor interactionJankMonitor) {
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler) {
         return new ExitDesktopTaskTransitionHandler(
-            transitions, context, interactionJankMonitor);
+            transitions, context, interactionJankMonitor, handler);
     }
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip1Module.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip1Module.java
index 037fbb235bd4..3a4764d45f2c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip1Module.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip1Module.java
@@ -101,7 +101,8 @@ public abstract class Pip1Module {
             DisplayInsetsController displayInsetsController,
             TabletopModeController pipTabletopController,
             Optional<OneHandedController> oneHandedController,
-            @ShellMainThread ShellExecutor mainExecutor) {
+            @ShellMainThread ShellExecutor mainExecutor,
+            @ShellMainThread Handler handler) {
         return Optional.ofNullable(PipController.create(
                 context, shellInit, shellCommandHandler, shellController,
                 displayController, pipAnimationController, pipAppOpsListener,
@@ -111,7 +112,7 @@ public abstract class Pip1Module {
                 pipTransitionState, pipTouchHandler, pipTransitionController,
                 windowManagerShellWrapper, taskStackListener, pipParamsChangedForwarder,
                 displayInsetsController, pipTabletopController, oneHandedController,
-                mainExecutor));
+                mainExecutor, handler));
     }
 
     // Handler is used by Icon.loadDrawableAsync
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index c74f4a7854d9..853284a58904 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -35,6 +35,7 @@ import android.graphics.PointF
 import android.graphics.Rect
 import android.graphics.Region
 import android.os.Binder
+import android.os.Handler
 import android.os.IBinder
 import android.os.SystemProperties
 import android.util.Size
@@ -136,7 +137,8 @@ class DesktopTasksController(
     @ShellMainThread private val mainExecutor: ShellExecutor,
     private val desktopTasksLimiter: Optional<DesktopTasksLimiter>,
     private val recentTasksController: RecentTasksController?,
-    private val interactionJankMonitor: InteractionJankMonitor
+    private val interactionJankMonitor: InteractionJankMonitor,
+    @ShellMainThread private val handler: Handler,
 ) :
     RemoteCallable<DesktopTasksController>,
     Transitions.TransitionHandler,
@@ -387,7 +389,7 @@ class DesktopTasksController(
         taskSurface: SurfaceControl,
     ) {
         logV("startDragToDesktop taskId=%d", taskInfo.taskId)
-        interactionJankMonitor.begin(taskSurface, context,
+        interactionJankMonitor.begin(taskSurface, context, handler,
             CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_HOLD)
         dragToDesktopTransitionHandler.startDragToDesktopTransition(
             taskInfo.taskId,
@@ -791,7 +793,7 @@ class DesktopTasksController(
         releaseVisualIndicator()
         if (!taskInfo.isResizeable && DesktopModeFlags.DISABLE_SNAP_RESIZE.isEnabled(context)) {
             interactionJankMonitor.begin(
-                taskSurface, context, CUJ_DESKTOP_MODE_SNAP_RESIZE, "drag_non_resizable"
+                taskSurface, context, handler, CUJ_DESKTOP_MODE_SNAP_RESIZE, "drag_non_resizable"
             )
 
             // reposition non-resizable app back to its original position before being dragged
@@ -804,7 +806,7 @@ class DesktopTasksController(
             )
         } else {
             interactionJankMonitor.begin(
-                taskSurface, context, CUJ_DESKTOP_MODE_SNAP_RESIZE, "drag_resizable"
+                taskSurface, context, handler, CUJ_DESKTOP_MODE_SNAP_RESIZE, "drag_resizable"
             )
             snapToHalfScreen(taskInfo, taskSurface, currentDragBounds, position)
         }
@@ -1604,7 +1606,7 @@ class DesktopTasksController(
         when (indicatorType) {
             IndicatorType.TO_DESKTOP_INDICATOR -> {
                 // Start a new jank interaction for the drag release to desktop window animation.
-                interactionJankMonitor.begin(taskSurface, context,
+                interactionJankMonitor.begin(taskSurface, context, handler,
                     CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_RELEASE, "to_desktop")
                 finalizeDragToDesktop(taskInfo)
             }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
index 597637d3fbfc..dae37f4926f5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
@@ -18,6 +18,7 @@ package com.android.wm.shell.desktopmode
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.content.Context
+import android.os.Handler
 import android.os.IBinder
 import android.view.SurfaceControl
 import android.view.WindowManager.TRANSIT_TO_BACK
@@ -29,6 +30,7 @@ import com.android.internal.jank.InteractionJankMonitor
 import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.protolog.ShellProtoLogGroup
+import com.android.wm.shell.shared.annotations.ShellMainThread
 import com.android.wm.shell.transition.Transitions
 import com.android.wm.shell.transition.Transitions.TransitionObserver
 
@@ -45,7 +47,8 @@ class DesktopTasksLimiter (
         private val shellTaskOrganizer: ShellTaskOrganizer,
         private val maxTasksLimit: Int,
         private val interactionJankMonitor: InteractionJankMonitor,
-        private val context: Context
+        private val context: Context,
+        @ShellMainThread private val handler: Handler,
 ) {
     private val minimizeTransitionObserver = MinimizeTransitionObserver()
     @VisibleForTesting
@@ -125,7 +128,7 @@ class DesktopTasksLimiter (
             if (mActiveTaskDetails != null && mActiveTaskDetails.transitionInfo != null) {
                 // Begin minimize window CUJ instrumentation.
                 interactionJankMonitor.begin(
-                    mActiveTaskDetails.transitionInfo?.rootLeash, context,
+                    mActiveTaskDetails.transitionInfo?.rootLeash, context, handler,
                     CUJ_DESKTOP_MODE_MINIMIZE_WINDOW
                 )
             }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandler.java
index e87be520cc91..dedd44f3950a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandler.java
@@ -29,6 +29,7 @@ import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.os.IBinder;
 import android.util.DisplayMetrics;
 import android.view.SurfaceControl;
@@ -44,6 +45,7 @@ import androidx.annotation.Nullable;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.jank.Cuj;
 import com.android.internal.jank.InteractionJankMonitor;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource;
 import com.android.wm.shell.transition.Transitions;
 
@@ -63,6 +65,8 @@ public class ExitDesktopTaskTransitionHandler implements Transitions.TransitionH
     private final Context mContext;
     private final Transitions mTransitions;
     private final InteractionJankMonitor mInteractionJankMonitor;
+    @ShellMainThread
+    private final Handler mHandler;
     private final List<IBinder> mPendingTransitionTokens = new ArrayList<>();
     private Consumer<SurfaceControl.Transaction> mOnAnimationFinishedCallback;
     private final Supplier<SurfaceControl.Transaction> mTransactionSupplier;
@@ -71,20 +75,24 @@ public class ExitDesktopTaskTransitionHandler implements Transitions.TransitionH
     public ExitDesktopTaskTransitionHandler(
             Transitions transitions,
             Context context,
-            InteractionJankMonitor interactionJankMonitor
-            ) {
-        this(transitions, SurfaceControl.Transaction::new, context, interactionJankMonitor);
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler
+    ) {
+        this(transitions, SurfaceControl.Transaction::new, context, interactionJankMonitor,
+                handler);
     }
 
     private ExitDesktopTaskTransitionHandler(
             Transitions transitions,
             Supplier<SurfaceControl.Transaction> supplier,
             Context context,
-            InteractionJankMonitor interactionJankMonitor) {
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler) {
         mTransitions = transitions;
         mTransactionSupplier = supplier;
         mContext = context;
         mInteractionJankMonitor = interactionJankMonitor;
+        mHandler = handler;
     }
 
     /**
@@ -154,7 +162,7 @@ public class ExitDesktopTaskTransitionHandler implements Transitions.TransitionH
             final SurfaceControl sc = change.getLeash();
             final Rect endBounds = change.getEndAbsBounds();
             mInteractionJankMonitor
-                .begin(sc, mContext, Cuj.CUJ_DESKTOP_MODE_EXIT_MODE);
+                    .begin(sc, mContext, mHandler, Cuj.CUJ_DESKTOP_MODE_EXIT_MODE);
             // Hide the first (fullscreen) frame because the animation will start from the freeform
             // size.
             startT.hide(sc)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
index 832e2d2bc77b..517e20910f6d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
@@ -28,6 +28,7 @@ import android.app.ActivityManager;
 import android.app.WindowConfiguration;
 import android.content.Context;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.os.IBinder;
 import android.util.ArrayMap;
 import android.view.SurfaceControl;
@@ -43,6 +44,7 @@ import com.android.internal.jank.InteractionJankMonitor;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
@@ -65,6 +67,8 @@ public class FreeformTaskTransitionHandler
     private final InteractionJankMonitor mInteractionJankMonitor;
     private final ShellExecutor mMainExecutor;
     private final ShellExecutor mAnimExecutor;
+    @ShellMainThread
+    private final Handler mHandler;
 
     private final List<IBinder> mPendingTransitionTokens = new ArrayList<>();
 
@@ -79,7 +83,8 @@ public class FreeformTaskTransitionHandler
             ShellExecutor mainExecutor,
             ShellExecutor animExecutor,
             DesktopModeTaskRepository desktopModeTaskRepository,
-            InteractionJankMonitor interactionJankMonitor) {
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler) {
         mTransitions = transitions;
         mContext = context;
         mWindowDecorViewModel = windowDecorViewModel;
@@ -88,6 +93,7 @@ public class FreeformTaskTransitionHandler
         mInteractionJankMonitor = interactionJankMonitor;
         mMainExecutor = mainExecutor;
         mAnimExecutor = animExecutor;
+        mHandler = handler;
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             shellInit.addInitCallback(this::onInit, this);
         }
@@ -267,7 +273,7 @@ public class FreeformTaskTransitionHandler
                         change.getTaskInfo().displayId) == 1) {
             // Starting the jank trace if closing the last window in desktop mode.
             mInteractionJankMonitor.begin(
-                    sc, mContext, CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE);
+                    sc, mContext, mHandler, CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE);
         }
         animator.addListener(
                 new AnimatorListenerAdapter() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java
index 1827923aad90..15472ebc149b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java
@@ -55,6 +55,7 @@ import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TaskStackListenerCallback;
 import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.shared.annotations.ExternalThread;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.KeyguardChangeListener;
 import com.android.wm.shell.sysui.ShellCommandHandler;
@@ -203,7 +204,7 @@ public class OneHandedController implements RemoteCallable<OneHandedController>,
             DisplayController displayController, DisplayLayout displayLayout,
             TaskStackListenerImpl taskStackListener,
             InteractionJankMonitor jankMonitor, UiEventLogger uiEventLogger,
-            ShellExecutor mainExecutor, Handler mainHandler) {
+            ShellExecutor mainExecutor, @ShellMainThread Handler mainHandler) {
         OneHandedSettingsUtil settingsUtil = new OneHandedSettingsUtil();
         OneHandedAccessibilityUtil accessibilityUtil = new OneHandedAccessibilityUtil(context);
         OneHandedTimeoutHandler timeoutHandler = new OneHandedTimeoutHandler(mainExecutor);
@@ -217,7 +218,7 @@ public class OneHandedController implements RemoteCallable<OneHandedController>,
                 mainExecutor);
         OneHandedDisplayAreaOrganizer organizer = new OneHandedDisplayAreaOrganizer(
                 context, displayLayout, settingsUtil, animationController, tutorialHandler,
-                jankMonitor, mainExecutor);
+                jankMonitor, mainExecutor, mainHandler);
         OneHandedUiEventLogger oneHandedUiEventsLogger = new OneHandedUiEventLogger(uiEventLogger);
         return new OneHandedController(context, shellInit, shellCommandHandler, shellController,
                 displayController, organizer, touchHandler, tutorialHandler, settingsUtil,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
index d157ca837608..95e633d0b5ec 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
@@ -23,6 +23,7 @@ import static com.android.wm.shell.onehanded.OneHandedAnimationController.TRANSI
 
 import android.content.Context;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.os.SystemProperties;
 import android.text.TextUtils;
 import android.util.ArrayMap;
@@ -42,6 +43,7 @@ import com.android.internal.jank.InteractionJankMonitor;
 import com.android.wm.shell.R;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.ShellExecutor;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -70,6 +72,8 @@ public class OneHandedDisplayAreaOrganizer extends DisplayAreaOrganizer {
     private final OneHandedSettingsUtil mOneHandedSettingsUtil;
     private final InteractionJankMonitor mJankMonitor;
     private final Context mContext;
+    @ShellMainThread
+    private final Handler mHandler;
 
     private boolean mIsReady;
     private float mLastVisualOffset = 0;
@@ -136,9 +140,11 @@ public class OneHandedDisplayAreaOrganizer extends DisplayAreaOrganizer {
             OneHandedAnimationController animationController,
             OneHandedTutorialHandler tutorialHandler,
             InteractionJankMonitor jankMonitor,
-            ShellExecutor mainExecutor) {
+            ShellExecutor mainExecutor,
+            @ShellMainThread Handler handler) {
         super(mainExecutor);
         mContext = context;
+        mHandler = handler;
         setDisplayLayout(displayLayout);
         mOneHandedSettingsUtil = oneHandedSettingsUtil;
         mAnimationController = animationController;
@@ -333,7 +339,7 @@ public class OneHandedDisplayAreaOrganizer extends DisplayAreaOrganizer {
                 getDisplayAreaTokenMap().entrySet().iterator().next();
         final InteractionJankMonitor.Configuration.Builder builder =
                 InteractionJankMonitor.Configuration.Builder.withSurface(
-                        cujType, mContext, firstEntry.getValue());
+                        cujType, mContext, firstEntry.getValue(), mHandler);
         if (!TextUtils.isEmpty(tag)) {
             builder.setTag(tag);
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
index deb7691f2d4d..af6844262771 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
@@ -44,6 +44,7 @@ import android.content.pm.ActivityInfo;
 import android.content.res.Configuration;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.os.SystemProperties;
 import android.util.Pair;
@@ -93,6 +94,7 @@ import com.android.wm.shell.pip.PipTaskOrganizer;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip.PipTransitionState;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.KeyguardChangeListener;
 import com.android.wm.shell.sysui.ShellCommandHandler;
@@ -146,6 +148,8 @@ public class PipController implements PipTransitionController.PipTransitionCallb
     private Optional<OneHandedController> mOneHandedController;
     private final ShellCommandHandler mShellCommandHandler;
     private final ShellController mShellController;
+    @ShellMainThread
+    private final Handler mHandler;
     protected final PipImpl mImpl;
 
     private final Rect mTmpInsetBounds = new Rect();
@@ -405,7 +409,8 @@ public class PipController implements PipTransitionController.PipTransitionCallb
             DisplayInsetsController displayInsetsController,
             TabletopModeController pipTabletopController,
             Optional<OneHandedController> oneHandedController,
-            ShellExecutor mainExecutor) {
+            ShellExecutor mainExecutor,
+            Handler handler) {
         if (!context.getPackageManager().hasSystemFeature(FEATURE_PICTURE_IN_PICTURE)) {
             ProtoLog.w(ShellProtoLogGroup.WM_SHELL_PICTURE_IN_PICTURE,
                     "%s: Device doesn't support Pip feature", TAG);
@@ -418,7 +423,8 @@ public class PipController implements PipTransitionController.PipTransitionCallb
                 pipDisplayLayoutState, pipMotionHelper, pipMediaController, phonePipMenuController,
                 pipTaskOrganizer, pipTransitionState, pipTouchHandler, pipTransitionController,
                 windowManagerShellWrapper, taskStackListener, pipParamsChangedForwarder,
-                displayInsetsController, pipTabletopController, oneHandedController, mainExecutor)
+                displayInsetsController, pipTabletopController, oneHandedController, mainExecutor,
+                handler)
                 .mImpl;
     }
 
@@ -446,11 +452,13 @@ public class PipController implements PipTransitionController.PipTransitionCallb
             DisplayInsetsController displayInsetsController,
             TabletopModeController tabletopModeController,
             Optional<OneHandedController> oneHandedController,
-            ShellExecutor mainExecutor
+            ShellExecutor mainExecutor,
+            @ShellMainThread Handler handler
     ) {
         mContext = context;
         mShellCommandHandler = shellCommandHandler;
         mShellController = shellController;
+        mHandler = handler;
         mImpl = new PipImpl();
         mWindowManagerShellWrapper = windowManagerShellWrapper;
         mDisplayController = displayController;
@@ -1047,7 +1055,8 @@ public class PipController implements PipTransitionController.PipTransitionCallb
         // Begin InteractionJankMonitor with PIP transition CUJs
         final InteractionJankMonitor.Configuration.Builder builder =
                 InteractionJankMonitor.Configuration.Builder.withSurface(
-                        CUJ_PIP_TRANSITION, mContext, mPipTaskOrganizer.getSurfaceControl())
+                                CUJ_PIP_TRANSITION, mContext, mPipTaskOrganizer.getSurfaceControl(),
+                                mHandler)
                 .setTag(getTransitionTag(direction))
                 .setTimeout(2000);
         InteractionJankMonitor.getInstance().begin(builder);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 4ba84eeb2e6d..5a905cfd317f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -1654,7 +1654,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             mSplitLayout = new SplitLayout(TAG + "SplitDivider", mContext,
                     mRootTaskInfo.configuration, this, mParentContainerCallbacks,
                     mDisplayController, mDisplayImeController, mTaskOrganizer,
-                    PARALLAX_ALIGN_CENTER /* parallaxType */);
+                    PARALLAX_ALIGN_CENTER /* parallaxType */, mMainHandler);
             mDisplayInsetsController.addInsetsChangedListener(mDisplayId, mSplitLayout);
         }
 
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 b14283f878a3..dc27c82d5ffa 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
@@ -111,6 +111,7 @@ import com.android.wm.shell.desktopmode.DesktopTasksLimiter;
 import com.android.wm.shell.desktopmode.DesktopWallpaperActivity;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.shared.desktopmode.DesktopModeFlags;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource;
@@ -154,7 +155,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
     private final ShellTaskOrganizer mTaskOrganizer;
     private final ShellController mShellController;
     private final Context mContext;
-    private final Handler mMainHandler;
+    private final @ShellMainThread Handler mMainHandler;
     private final @ShellBackgroundThread ShellExecutor mBgExecutor;
     private final Choreographer mMainChoreographer;
     private final DisplayController mDisplayController;
@@ -214,7 +215,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
     public DesktopModeWindowDecorViewModel(
             Context context,
             ShellExecutor shellExecutor,
-            Handler mainHandler,
+            @ShellMainThread Handler mainHandler,
             Choreographer mainChoreographer,
             @ShellBackgroundThread ShellExecutor bgExecutor,
             ShellInit shellInit,
@@ -270,7 +271,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
     DesktopModeWindowDecorViewModel(
             Context context,
             ShellExecutor shellExecutor,
-            Handler mainHandler,
+            @ShellMainThread Handler mainHandler,
             Choreographer mainChoreographer,
             @ShellBackgroundThread ShellExecutor bgExecutor,
             ShellInit shellInit,
@@ -495,7 +496,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
             return;
         }
         mInteractionJankMonitor.begin(
-                decoration.mTaskSurface, mContext, Cuj.CUJ_DESKTOP_MODE_MAXIMIZE_WINDOW, source);
+                decoration.mTaskSurface, mContext, mMainHandler,
+                Cuj.CUJ_DESKTOP_MODE_MAXIMIZE_WINDOW, source);
         mDesktopTasksController.toggleDesktopTaskSize(decoration.mTaskInfo);
         decoration.closeHandleMenu();
         decoration.closeMaximizeMenu();
@@ -512,7 +514,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
             Toast.makeText(mContext,
                     R.string.desktop_mode_non_resizable_snap_text, Toast.LENGTH_SHORT).show();
         } else {
-            mInteractionJankMonitor.begin(decoration.mTaskSurface, mContext,
+            mInteractionJankMonitor.begin(decoration.mTaskSurface, mContext, mMainHandler,
                     Cuj.CUJ_DESKTOP_MODE_SNAP_RESIZE, "maximize_menu_resizable");
             mDesktopTasksController.snapToHalfScreen(
                     decoration.mTaskInfo,
@@ -548,7 +550,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
             return;
         }
         final WindowContainerTransaction wct = new WindowContainerTransaction();
-        mInteractionJankMonitor.begin(decoration.mTaskSurface, mContext,
+        mInteractionJankMonitor.begin(decoration.mTaskSurface, mContext, mMainHandler,
                 CUJ_DESKTOP_MODE_ENTER_MODE_APP_HANDLE_MENU);
         // App sometimes draws before the insets from WindowDecoration#relayout have
         // been added, so they must be added here
@@ -1378,7 +1380,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
                 mDragStartListener,
                 mTransitions,
                 mInteractionJankMonitor,
-                mTransactionFactory);
+                mTransactionFactory,
+                mMainHandler);
         windowDecoration.setTaskDragResizer(taskPositioner);
 
         final DesktopModeTouchEventListener touchEventListener =
@@ -1602,7 +1605,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
                 DragPositioningCallbackUtility.DragStartListener dragStartListener,
                 Transitions transitions,
                 InteractionJankMonitor interactionJankMonitor,
-                Supplier<SurfaceControl.Transaction> transactionFactory) {
+                Supplier<SurfaceControl.Transaction> transactionFactory,
+                Handler handler) {
             final TaskPositioner taskPositioner = DesktopModeStatus.isVeiledResizeEnabled()
                     ? new VeiledResizeTaskPositioner(
                             taskOrganizer,
@@ -1610,7 +1614,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
                             displayController,
                             dragStartListener,
                             transitions,
-                            interactionJankMonitor)
+                            interactionJankMonitor,
+                            handler)
                     : new FluidResizeTaskPositioner(
                             taskOrganizer,
                             transitions,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositioner.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositioner.java
index 599815530f63..6f3f41191485 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositioner.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositioner.java
@@ -24,6 +24,7 @@ import static com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_RESIZE_WINDOW;
 import android.graphics.Point;
 import android.graphics.PointF;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.os.IBinder;
 import android.view.Surface;
 import android.view.SurfaceControl;
@@ -37,6 +38,7 @@ import androidx.annotation.Nullable;
 import com.android.internal.jank.InteractionJankMonitor;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
+import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.transition.Transitions;
 
 import java.util.function.Supplier;
@@ -63,14 +65,17 @@ public class VeiledResizeTaskPositioner implements TaskPositioner, Transitions.T
     private int mCtrlType;
     private boolean mIsResizingOrAnimatingResize;
     @Surface.Rotation private int mRotation;
+    @ShellMainThread
+    private final Handler mHandler;
 
     public VeiledResizeTaskPositioner(ShellTaskOrganizer taskOrganizer,
             DesktopModeWindowDecoration windowDecoration,
             DisplayController displayController,
             DragPositioningCallbackUtility.DragStartListener dragStartListener,
-            Transitions transitions, InteractionJankMonitor interactionJankMonitor) {
+            Transitions transitions, InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler) {
         this(taskOrganizer, windowDecoration, displayController, dragStartListener,
-                SurfaceControl.Transaction::new, transitions, interactionJankMonitor);
+                SurfaceControl.Transaction::new, transitions, interactionJankMonitor, handler);
     }
 
     public VeiledResizeTaskPositioner(ShellTaskOrganizer taskOrganizer,
@@ -78,7 +83,7 @@ public class VeiledResizeTaskPositioner implements TaskPositioner, Transitions.T
             DisplayController displayController,
             DragPositioningCallbackUtility.DragStartListener dragStartListener,
             Supplier<SurfaceControl.Transaction> supplier, Transitions transitions,
-            InteractionJankMonitor interactionJankMonitor) {
+            InteractionJankMonitor interactionJankMonitor, @ShellMainThread Handler handler) {
         mDesktopWindowDecoration = windowDecoration;
         mTaskOrganizer = taskOrganizer;
         mDisplayController = displayController;
@@ -86,6 +91,7 @@ public class VeiledResizeTaskPositioner implements TaskPositioner, Transitions.T
         mTransactionSupplier = supplier;
         mTransitions = transitions;
         mInteractionJankMonitor = interactionJankMonitor;
+        mHandler = handler;
     }
 
     @Override
@@ -97,7 +103,7 @@ public class VeiledResizeTaskPositioner implements TaskPositioner, Transitions.T
         if (isResizing()) {
             // Capture CUJ for re-sizing window in DW mode.
             mInteractionJankMonitor.begin(mDesktopWindowDecoration.mTaskSurface,
-                    mDesktopWindowDecoration.mContext, CUJ_DESKTOP_MODE_RESIZE_WINDOW);
+                    mDesktopWindowDecoration.mContext, mHandler, CUJ_DESKTOP_MODE_RESIZE_WINDOW);
             if (!mDesktopWindowDecoration.mTaskInfo.isFocused) {
                 WindowContainerTransaction wct = new WindowContainerTransaction();
                 wct.reorder(mDesktopWindowDecoration.mTaskInfo.token, true);
@@ -131,7 +137,7 @@ public class VeiledResizeTaskPositioner implements TaskPositioner, Transitions.T
         } else if (mCtrlType == CTRL_TYPE_UNDEFINED) {
             // Begin window drag CUJ instrumentation only when drag position moves.
             mInteractionJankMonitor.begin(mDesktopWindowDecoration.mTaskSurface,
-                    mDesktopWindowDecoration.mContext, CUJ_DESKTOP_MODE_DRAG_WINDOW);
+                    mDesktopWindowDecoration.mContext, mHandler, CUJ_DESKTOP_MODE_DRAG_WINDOW);
             final SurfaceControl.Transaction t = mTransactionSupplier.get();
             DragPositioningCallbackUtility.setPositionOnDrag(mDesktopWindowDecoration,
                     mRepositionTaskBounds, mTaskBoundsAtDragStart, mRepositionStartPoint, t, x, y);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
index b53ea3837178..227060d15640 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
@@ -137,6 +137,8 @@ public class BackAnimationControllerTest extends ShellTestCase {
     private Transitions mTransitions;
     @Mock
     private RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
+    @Mock
+    private Handler mHandler;
 
     private BackAnimationController mController;
     private TestableContentResolver mContentResolver;
@@ -161,13 +163,14 @@ public class BackAnimationControllerTest extends ShellTestCase {
         mTestableLooper = TestableLooper.get(this);
         mShellInit = spy(new ShellInit(mShellExecutor));
         mDefaultCrossActivityBackAnimation = new DefaultCrossActivityBackAnimation(mContext,
-                mAnimationBackground, mRootTaskDisplayAreaOrganizer);
-        mCrossTaskBackAnimation = new CrossTaskBackAnimation(mContext, mAnimationBackground);
+                mAnimationBackground, mRootTaskDisplayAreaOrganizer, mHandler);
+        mCrossTaskBackAnimation = new CrossTaskBackAnimation(mContext, mAnimationBackground,
+                mHandler);
         mShellBackAnimationRegistry =
                 new ShellBackAnimationRegistry(mDefaultCrossActivityBackAnimation,
                         mCrossTaskBackAnimation, /* dialogCloseAnimation= */ null,
                         new CustomCrossActivityBackAnimation(mContext, mAnimationBackground,
-                                mRootTaskDisplayAreaOrganizer),
+                                mRootTaskDisplayAreaOrganizer, mHandler),
                         /* defaultBackToHomeAnimation= */ null);
         mController =
                 new BackAnimationController(
@@ -181,7 +184,8 @@ public class BackAnimationControllerTest extends ShellTestCase {
                         mAnimationBackground,
                         mShellBackAnimationRegistry,
                         mShellCommandHandler,
-                        mTransitions);
+                        mTransitions,
+                        mHandler);
         mShellInit.init();
         mShellExecutor.flushAll();
         mTouchableRegion = new Rect(0, 0, 100, 100);
@@ -344,7 +348,8 @@ public class BackAnimationControllerTest extends ShellTestCase {
                         mAnimationBackground,
                         mShellBackAnimationRegistry,
                         mShellCommandHandler,
-                        mTransitions);
+                        mTransitions,
+                        mHandler);
         shellInit.init();
         registerAnimation(BackNavigationInfo.TYPE_RETURN_TO_HOME);
 
@@ -898,7 +903,8 @@ public class BackAnimationControllerTest extends ShellTestCase {
                 new BackAnimationRunner(
                         mAnimatorCallback,
                         mBackAnimationRunner,
-                        mContext));
+                        mContext,
+                        mHandler));
     }
 
     private void unregisterAnimation(int type) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/CustomCrossActivityBackAnimationTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/CustomCrossActivityBackAnimationTest.kt
index 080ad901c656..5b5ef6f48789 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/CustomCrossActivityBackAnimationTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/CustomCrossActivityBackAnimationTest.kt
@@ -22,6 +22,7 @@ import android.app.WindowConfiguration
 import android.graphics.Color
 import android.graphics.Point
 import android.graphics.Rect
+import android.os.Handler
 import android.os.RemoteException
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
@@ -66,6 +67,7 @@ class CustomCrossActivityBackAnimationTest : ShellTestCase() {
     @Mock private lateinit var transitionAnimation: TransitionAnimation
     @Mock private lateinit var appCompatTaskInfo: AppCompatTaskInfo
     @Mock private lateinit var transaction: Transaction
+    @Mock private lateinit var handler: Handler
 
     private lateinit var customCrossActivityBackAnimation: CustomCrossActivityBackAnimation
     private lateinit var customAnimationLoader: CustomAnimationLoader
@@ -80,7 +82,8 @@ class CustomCrossActivityBackAnimationTest : ShellTestCase() {
                 backAnimationBackground,
                 rootTaskDisplayAreaOrganizer,
                 transaction,
-                customAnimationLoader
+                customAnimationLoader,
+                handler,
             )
 
         whenever(transitionAnimation.loadAppTransitionAnimation(eq(PACKAGE_NAME), eq(OPEN_RES_ID)))
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/DividerViewTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/DividerViewTest.java
index f5847cc27071..cf69704a0470 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/DividerViewTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/DividerViewTest.java
@@ -25,6 +25,7 @@ import static org.mockito.Mockito.verify;
 
 import android.content.res.Configuration;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.os.SystemClock;
 import android.provider.DeviceConfig;
 import android.view.InputDevice;
@@ -55,6 +56,7 @@ public class DividerViewTest extends ShellTestCase {
     private @Mock DisplayController mDisplayController;
     private @Mock DisplayImeController mDisplayImeController;
     private @Mock ShellTaskOrganizer mTaskOrganizer;
+    private @Mock Handler mHandler;
     private SplitLayout mSplitLayout;
     private DividerView mDividerView;
 
@@ -65,7 +67,7 @@ public class DividerViewTest extends ShellTestCase {
         Configuration configuration = getConfiguration();
         mSplitLayout = new SplitLayout("TestSplitLayout", mContext, configuration,
                 mSplitLayoutHandler, mCallbacks, mDisplayController, mDisplayImeController,
-                mTaskOrganizer, SplitLayout.PARALLAX_NONE);
+                mTaskOrganizer, SplitLayout.PARALLAX_NONE, mHandler);
         SplitWindowManager splitWindowManager = new SplitWindowManager("TestSplitWindowManager",
                 mContext,
                 configuration, mCallbacks);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java
index 82b3a7de521b..177e47a342f6 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java
@@ -35,6 +35,7 @@ import static org.mockito.Mockito.verify;
 import android.app.ActivityManager;
 import android.content.res.Configuration;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.window.WindowContainerTransaction;
 
 import androidx.test.annotation.UiThreadTest;
@@ -65,6 +66,7 @@ public class SplitLayoutTests extends ShellTestCase {
     @Mock DisplayImeController mDisplayImeController;
     @Mock ShellTaskOrganizer mTaskOrganizer;
     @Mock WindowContainerTransaction mWct;
+    @Mock Handler mHandler;
     @Captor ArgumentCaptor<Runnable> mRunnableCaptor;
     private SplitLayout mSplitLayout;
 
@@ -80,7 +82,8 @@ public class SplitLayoutTests extends ShellTestCase {
                 mDisplayController,
                 mDisplayImeController,
                 mTaskOrganizer,
-                SplitLayout.PARALLAX_NONE));
+                SplitLayout.PARALLAX_NONE,
+                mHandler));
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 5ef34a5518e2..712c30641200 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -38,6 +38,7 @@ import android.graphics.Point
 import android.graphics.PointF
 import android.graphics.Rect
 import android.os.Binder
+import android.os.Handler
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
@@ -177,6 +178,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
   private lateinit var mockInteractionJankMonitor: InteractionJankMonitor
   @Mock private lateinit var mockSurface: SurfaceControl
   @Mock private lateinit var taskbarDesktopTaskListener: TaskbarDesktopTaskListener
+  @Mock private lateinit var mockHandler: Handler
 
   private lateinit var mockitoSession: StaticMockitoSession
   private lateinit var controller: DesktopTasksController
@@ -217,7 +219,8 @@ class DesktopTasksControllerTest : ShellTestCase() {
             shellTaskOrganizer,
             MAX_TASK_LIMIT,
             mockInteractionJankMonitor,
-            mContext)
+            mContext,
+            mockHandler)
 
     whenever(shellTaskOrganizer.getRunningTasks(anyInt())).thenAnswer { runningTasks }
     whenever(transitions.startTransition(anyInt(), any(), isNull())).thenAnswer { Binder() }
@@ -270,7 +273,9 @@ class DesktopTasksControllerTest : ShellTestCase() {
         shellExecutor,
         Optional.of(desktopTasksLimiter),
         recentTasksController,
-        mockInteractionJankMonitor)
+        mockInteractionJankMonitor,
+        mockHandler,
+      )
   }
 
   @After
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
index 2d0e428c45cb..61d03cac035c 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
@@ -18,6 +18,7 @@ package com.android.wm.shell.desktopmode
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.os.Binder
+import android.os.Handler
 import android.platform.test.flag.junit.SetFlagsRule
 import android.testing.AndroidTestingRunner
 import android.view.Display.DEFAULT_DISPLAY
@@ -70,6 +71,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
     @Mock lateinit var shellTaskOrganizer: ShellTaskOrganizer
     @Mock lateinit var transitions: Transitions
     @Mock lateinit var interactionJankMonitor: InteractionJankMonitor
+    @Mock lateinit var handler: Handler
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var desktopTasksLimiter: DesktopTasksLimiter
@@ -85,7 +87,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
 
         desktopTasksLimiter =
             DesktopTasksLimiter(transitions, desktopTaskRepo, shellTaskOrganizer, MAX_TASK_LIMIT,
-                interactionJankMonitor, mContext)
+                interactionJankMonitor, mContext, handler)
     }
 
     @After
@@ -97,7 +99,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
     fun createDesktopTasksLimiter_withZeroLimit_shouldThrow() {
         assertFailsWith<IllegalArgumentException> {
             DesktopTasksLimiter(transitions, desktopTaskRepo, shellTaskOrganizer, 0,
-                interactionJankMonitor, mContext)
+                interactionJankMonitor, mContext, handler)
         }
     }
 
@@ -105,7 +107,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
     fun createDesktopTasksLimiter_withNegativeLimit_shouldThrow() {
         assertFailsWith<IllegalArgumentException> {
             DesktopTasksLimiter(transitions, desktopTaskRepo, shellTaskOrganizer, -5,
-                interactionJankMonitor, mContext)
+                interactionJankMonitor, mContext, handler)
         }
     }
 
@@ -334,7 +336,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
     fun getTaskToMinimizeIfNeeded_tasksAboveLimit_otherLimit_returnsBackTask() {
         desktopTasksLimiter =
             DesktopTasksLimiter(transitions, desktopTaskRepo, shellTaskOrganizer, MAX_TASK_LIMIT2,
-                interactionJankMonitor, mContext)
+                interactionJankMonitor, mContext, handler)
         val tasks = (1..MAX_TASK_LIMIT2 + 1).map { setUpFreeformTask() }
 
         val minimizedTask = desktopTasksLimiter.getTaskToMinimizeIfNeeded(
@@ -375,6 +377,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
         verify(interactionJankMonitor).begin(
             any(),
             eq(mContext),
+            eq(handler),
             eq(CUJ_DESKTOP_MODE_MINIMIZE_WINDOW))
 
         desktopTasksLimiter.getTransitionObserver().onTransitionFinished(
@@ -403,7 +406,9 @@ class DesktopTasksLimiterTest : ShellTestCase() {
         verify(interactionJankMonitor).begin(
             any(),
             eq(mContext),
-            eq(CUJ_DESKTOP_MODE_MINIMIZE_WINDOW))
+            eq(handler),
+            eq(CUJ_DESKTOP_MODE_MINIMIZE_WINDOW),
+        )
 
         desktopTasksLimiter.getTransitionObserver().onTransitionFinished(
             transition,
@@ -432,6 +437,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
         verify(interactionJankMonitor).begin(
             any(),
             eq(mContext),
+            eq(handler),
             eq(CUJ_DESKTOP_MODE_MINIMIZE_WINDOW))
 
         desktopTasksLimiter.getTransitionObserver().onTransitionMerged(
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandlerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandlerTest.java
index e0463b41ad20..fefa933c5208 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandlerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/ExitDesktopTaskTransitionHandlerTest.java
@@ -34,6 +34,7 @@ import android.app.WindowConfiguration;
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Point;
+import android.os.Handler;
 import android.os.IBinder;
 import android.util.DisplayMetrics;
 import android.view.SurfaceControl;
@@ -81,6 +82,8 @@ public class ExitDesktopTaskTransitionHandlerTest extends ShellTestCase {
     Transitions.TransitionFinishCallback mTransitionFinishCallback;
     @Mock
     ShellExecutor mExecutor;
+    @Mock
+    Handler mHandler;
 
     private Point mPoint;
     private ExitDesktopTaskTransitionHandler mExitDesktopTaskTransitionHandler;
@@ -97,7 +100,7 @@ public class ExitDesktopTaskTransitionHandlerTest extends ShellTestCase {
                 .thenReturn(getContext().getResources().getDisplayMetrics());
 
         mExitDesktopTaskTransitionHandler = new ExitDesktopTaskTransitionHandler(mTransitions,
-                mContext, mInteractionJankMonitor);
+                mContext, mInteractionJankMonitor, mHandler);
         mPoint = new Point(0, 0);
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java
index 9c7f7237871a..9146906b6385 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java
@@ -37,6 +37,7 @@ import static org.mockito.Mockito.when;
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.Binder;
+import android.os.Handler;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.Display;
@@ -100,6 +101,8 @@ public class OneHandedDisplayAreaOrganizerTest extends OneHandedTestCase {
     OneHandedSettingsUtil mMockSettingsUitl;
     @Mock
     InteractionJankMonitor mJankMonitor;
+    @Mock
+    Handler mMockHandler;
 
     List<DisplayAreaAppearedInfo> mDisplayAreaAppearedInfoList = new ArrayList<>();
 
@@ -142,7 +145,8 @@ public class OneHandedDisplayAreaOrganizerTest extends OneHandedTestCase {
                 mMockAnimationController,
                 mTutorialHandler,
                 mJankMonitor,
-                mMockShellMainExecutor));
+                mMockShellMainExecutor,
+                mMockHandler));
 
         for (int i = 0; i < DISPLAYAREA_INFO_COUNT; i++) {
             mDisplayAreaAppearedInfoList.add(getDummyDisplayAreaInfo());
@@ -429,7 +433,8 @@ public class OneHandedDisplayAreaOrganizerTest extends OneHandedTestCase {
                         mMockAnimationController,
                         mTutorialHandler,
                         mJankMonitor,
-                        mMockShellMainExecutor));
+                        mMockShellMainExecutor,
+                        mMockHandler));
 
         assertThat(testSpiedDisplayAreaOrganizer.isReady()).isFalse();
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
index f3944d5ac352..96003515a485 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
@@ -40,6 +40,7 @@ import android.content.res.Configuration;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.Bundle;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
@@ -115,6 +116,7 @@ public class PipControllerTest extends ShellTestCase {
     @Mock private PipParamsChangedForwarder mMockPipParamsChangedForwarder;
     @Mock private DisplayInsetsController mMockDisplayInsetsController;
     @Mock private TabletopModeController mMockTabletopModeController;
+    @Mock private Handler mMockHandler;
 
     @Mock private DisplayLayout mMockDisplayLayout1;
     @Mock private DisplayLayout mMockDisplayLayout2;
@@ -138,7 +140,7 @@ public class PipControllerTest extends ShellTestCase {
                 mMockPipTransitionController, mMockWindowManagerShellWrapper,
                 mMockTaskStackListener, mMockPipParamsChangedForwarder,
                 mMockDisplayInsetsController, mMockTabletopModeController,
-                mMockOneHandedController, mMockExecutor);
+                mMockOneHandedController, mMockExecutor, mMockHandler);
         mShellInit.init();
         when(mMockPipBoundsAlgorithm.getSnapAlgorithm()).thenReturn(mMockPipSnapAlgorithm);
         when(mMockPipTouchHandler.getMotionHelper()).thenReturn(mMockPipMotionHelper);
@@ -230,7 +232,7 @@ public class PipControllerTest extends ShellTestCase {
                 mMockPipTransitionController, mMockWindowManagerShellWrapper,
                 mMockTaskStackListener, mMockPipParamsChangedForwarder,
                 mMockDisplayInsetsController, mMockTabletopModeController,
-                mMockOneHandedController, mMockExecutor));
+                mMockOneHandedController, mMockExecutor, mMockHandler));
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
index a18fbf0891ef..85bc7cc287e6 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
@@ -247,7 +247,18 @@ class DesktopModeWindowDecorViewModelTests : ShellTestCase() {
         whenever(mockDisplayController.getDisplayLayout(any())).thenReturn(mockDisplayLayout)
         whenever(mockDisplayLayout.stableInsets()).thenReturn(STABLE_INSETS)
         whenever(mockInputMonitorFactory.create(any(), any())).thenReturn(mockInputMonitor)
-        whenever(mockTaskPositionerFactory.create(any(), any(), any(), any(), any(), any(), any()))
+        whenever(
+            mockTaskPositionerFactory.create(
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any()
+            )
+        )
             .thenReturn(mockTaskPositioner)
 
         doReturn(mockToast).`when` { Toast.makeText(any(), anyInt(), anyInt()) }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositionerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositionerTest.kt
index 7784af6b1111..ab41d9c80177 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositionerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/VeiledResizeTaskPositionerTest.kt
@@ -21,6 +21,7 @@ import android.content.Context
 import android.content.res.Resources
 import android.graphics.Point
 import android.graphics.Rect
+import android.os.Handler
 import android.os.IBinder
 import android.testing.AndroidTestingRunner
 import android.view.Display
@@ -107,6 +108,8 @@ class VeiledResizeTaskPositionerTest : ShellTestCase() {
     private lateinit var mockResources: Resources
     @Mock
     private lateinit var mockInteractionJankMonitor: InteractionJankMonitor
+    @Mock
+    private lateinit var mockHandler: Handler
 
     private lateinit var taskPositioner: VeiledResizeTaskPositioner
 
@@ -155,7 +158,8 @@ class VeiledResizeTaskPositionerTest : ShellTestCase() {
                         mockDragStartListener,
                         mockTransactionFactory,
                         mockTransitions,
-                        mockInteractionJankMonitor
+                        mockInteractionJankMonitor,
+                        mockHandler,
                 )
     }