diff --git a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
index 6916a1d..e160f82 100644
--- a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
+++ b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
@@ -20,6 +20,7 @@
 import android.animation.AnimatorSet
 import android.animation.ValueAnimator
 import android.content.Context
+import android.graphics.Rect
 import android.os.IBinder
 import android.view.SurfaceControl.Transaction
 import android.view.WindowManager.TRANSIT_OPEN
@@ -31,6 +32,7 @@
 import android.window.TransitionInfo.Change
 import androidx.core.animation.addListener
 import com.android.app.animation.Interpolators
+import com.android.internal.policy.ScreenDecorationsUtils
 import com.android.quickstep.RemoteRunnable
 import com.android.wm.shell.shared.animation.MinimizeAnimator
 import com.android.wm.shell.shared.animation.WindowAnimator
@@ -43,8 +45,19 @@
  * ([android.view.WindowManager.TRANSIT_TO_BACK]) this transition will apply a minimize animation to
  * that window.
  */
-class DesktopAppLaunchTransition(private val context: Context, private val mainExecutor: Executor) :
-    RemoteTransitionStub() {
+class DesktopAppLaunchTransition(
+    private val context: Context,
+    private val mainExecutor: Executor,
+    private val launchType: AppLaunchType,
+) : RemoteTransitionStub() {
+
+    enum class AppLaunchType(
+        val boundsAnimationParams: WindowAnimator.BoundsAnimationParams,
+        val alphaDurationMs: Long,
+    ) {
+        LAUNCH(launchBoundsAnimationDef, /* alphaDurationMs= */ 200L),
+        UNMINIMIZE(unminimizeBoundsAnimationDef, /* alphaDurationMs= */ 100L),
+    }
 
     override fun startAnimation(
         token: IBinder,
@@ -105,18 +118,24 @@
         val boundsAnimator =
             WindowAnimator.createBoundsAnimator(
                 context.resources.displayMetrics,
-                launchBoundsAnimationDef,
+                launchType.boundsAnimationParams,
                 change,
                 transaction,
             )
         val alphaAnimator =
             ValueAnimator.ofFloat(0f, 1f).apply {
-                duration = LAUNCH_ANIM_ALPHA_DURATION_MS
+                duration = launchType.alphaDurationMs
                 interpolator = Interpolators.LINEAR
                 addUpdateListener { animation ->
                     transaction.setAlpha(change.leash, animation.animatedValue as Float).apply()
                 }
             }
+        val clipRect = Rect(change.endAbsBounds).apply { offsetTo(0, 0) }
+        transaction.setCrop(change.leash, clipRect)
+        transaction.setCornerRadius(
+            change.leash,
+            ScreenDecorationsUtils.getWindowCornerRadius(context),
+        )
         return AnimatorSet().apply {
             playTogether(boundsAnimator, alphaAnimator)
             addListener(onEnd = { animation -> onAnimFinish(animation) })
@@ -124,13 +143,18 @@
     }
 
     companion object {
-        private val LAUNCH_CHANGE_MODES = intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT)
-
-        private const val LAUNCH_ANIM_ALPHA_DURATION_MS = 100L
-        private const val MINIMIZE_ANIM_ALPHA_DURATION_MS = 100L
+        val LAUNCH_CHANGE_MODES = intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT)
 
         private val launchBoundsAnimationDef =
             WindowAnimator.BoundsAnimationParams(
+                durationMs = 600,
+                startOffsetYDp = 36f,
+                startScale = 0.95f,
+                interpolator = Interpolators.STANDARD_DECELERATE,
+            )
+
+        private val unminimizeBoundsAnimationDef =
+            WindowAnimator.BoundsAnimationParams(
                 durationMs = 300,
                 startOffsetYDp = 12f,
                 startScale = 0.97f,
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
index 5e11601..390112e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
@@ -36,6 +36,7 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.desktop.DesktopAppLaunchTransition;
+import com.android.launcher3.desktop.DesktopAppLaunchTransition.AppLaunchType;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayContext;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayDragLayer;
 import com.android.launcher3.util.DisplayController;
@@ -251,7 +252,8 @@
         ) {
             // This app is being unminimized - use our own transition runner.
             remoteTransition = new RemoteTransition(
-                    new DesktopAppLaunchTransition(context, MAIN_EXECUTOR));
+                    new DesktopAppLaunchTransition(
+                        context, MAIN_EXECUTOR, AppLaunchType.UNMINIMIZE));
         }
         mControllers.taskbarActivityContext.handleGroupTaskLaunch(
                 task,
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index c9c52b5..82acc0c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -59,6 +59,7 @@
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.hardware.display.DisplayManager;
+import android.os.Bundle;
 import android.os.IRemoteCallback;
 import android.os.Process;
 import android.os.Trace;
@@ -91,6 +92,7 @@
 import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.desktop.DesktopAppLaunchTransition;
+import com.android.launcher3.desktop.DesktopAppLaunchTransition.AppLaunchType;
 import com.android.launcher3.dot.DotInfo;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.folder.FolderIcon;
@@ -284,6 +286,7 @@
         BubbleBarController.onTaskbarRecreated();
         if (BubbleBarController.isBubbleBarEnabled()
                 && !mDeviceProfile.isPhone
+                && !mDeviceProfile.isVerticalBarLayout()
                 && bubbleBarView != null
         ) {
             Optional<BubbleStashedHandleViewController> bubbleHandleController = Optional.empty();
@@ -861,6 +864,33 @@
         return makeDefaultActivityOptions(SPLASH_SCREEN_STYLE_UNDEFINED);
     }
 
+    private ActivityOptionsWrapper getActivityLaunchDesktopOptions(ItemInfo info) {
+        if (!DesktopModeFlags.ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS.isTrue()) {
+            return null;
+        }
+        if (!areDesktopTasksVisible()) {
+            return null;
+        }
+        BubbleTextView.RunningAppState appState =
+                mControllers.taskbarRecentAppsController.getDesktopItemState(info);
+        AppLaunchType launchType = null;
+        switch (appState) {
+            case RUNNING:
+                return null;
+            case MINIMIZED:
+                launchType = AppLaunchType.UNMINIMIZE;
+                break;
+            case NOT_RUNNING:
+                launchType = AppLaunchType.LAUNCH;
+                break;
+        }
+        ActivityOptions options = ActivityOptions.makeRemoteTransition(
+                new RemoteTransition(
+                        new DesktopAppLaunchTransition(
+                                /* context= */ this, getMainExecutor(), launchType)));
+        return new ActivityOptionsWrapper(options, new RunnableList());
+    }
+
     /**
      * Sets a new data-source for this taskbar instance
      */
@@ -1401,7 +1431,9 @@
     }
 
     private RemoteTransition createUnminimizeRemoteTransition() {
-        return new RemoteTransition(new DesktopAppLaunchTransition(this, getMainExecutor()));
+        return new RemoteTransition(
+                new DesktopAppLaunchTransition(
+                        this, getMainExecutor(), AppLaunchType.UNMINIMIZE));
     }
 
     /**
@@ -1502,25 +1534,31 @@
                 .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
         try {
             TestLogging.recordEvent(TestProtocol.SEQUENCE_MAIN, "start: taskbarAppIcon");
-            if (info.user.equals(Process.myUserHandle())) {
-                // TODO(b/216683257): Use startActivityForResult for search results that require it.
-                if (taskInRecents != null) {
-                    // Re launch instance from recents
-                    ActivityOptionsWrapper opts = getActivityLaunchOptions(null, info);
-                    opts.options.setLaunchDisplayId(
-                            getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
-                    if (ActivityManagerWrapper.getInstance()
-                            .startActivityFromRecents(taskInRecents.key, opts.options)) {
-                        mControllers.uiController.getRecentsView()
-                                .addSideTaskLaunchCallback(opts.onEndCallback);
-                        return;
-                    }
-                }
-                startActivity(intent);
-            } else {
+            if (!info.user.equals(Process.myUserHandle())) {
+                // TODO b/376819104: support Desktop launch animations for apps in managed profiles
                 getSystemService(LauncherApps.class).startMainActivity(
                         intent.getComponent(), info.user, intent.getSourceBounds(), null);
+                return;
             }
+            // TODO(b/216683257): Use startActivityForResult for search results that require it.
+            if (taskInRecents != null) {
+                // Re launch instance from recents
+                ActivityOptionsWrapper opts = getActivityLaunchOptions(null, info);
+                opts.options.setLaunchDisplayId(
+                        getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
+                if (ActivityManagerWrapper.getInstance()
+                        .startActivityFromRecents(taskInRecents.key, opts.options)) {
+                    mControllers.uiController.getRecentsView()
+                            .addSideTaskLaunchCallback(opts.onEndCallback);
+                    return;
+                }
+            }
+            ActivityOptionsWrapper opts = null;
+            if (areDesktopTasksVisible()) {
+                opts = getActivityLaunchDesktopOptions(info);
+            }
+            Bundle optionsBundle = opts == null ? null : opts.options.toBundle();
+            startActivity(intent, optionsBundle);
         } catch (NullPointerException | ActivityNotFoundException | SecurityException e) {
             Toast.makeText(this, R.string.activity_not_found, Toast.LENGTH_SHORT)
                     .show();
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index 7b05043..3d57de4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -73,6 +73,33 @@
     var shownTasks: List<GroupTask> = emptyList()
         private set
 
+    /**
+     * Returns the state of the most active Desktop task represented by the given [ItemInfo].
+     *
+     * If there are several tasks represented by the same [ItemInfo] we return the most active one,
+     * i.e. we return [DesktopAppState.RUNNING] over [DesktopAppState.MINIMIZED], and
+     * [DesktopAppState.MINIMIZED] over [DesktopAppState.NOT_RUNNING].
+     */
+    fun getDesktopItemState(itemInfo: ItemInfo?): RunningAppState {
+        val packageName = itemInfo?.getTargetPackage() ?: return RunningAppState.NOT_RUNNING
+        return getDesktopAppState(packageName, itemInfo.user.identifier)
+    }
+
+    private fun getDesktopAppState(packageName: String, userId: Int): RunningAppState {
+        val tasks = desktopTask?.tasks ?: return RunningAppState.NOT_RUNNING
+        val appTasks =
+            tasks.filter { task ->
+                packageName == task.key.packageName && task.key.userId == userId
+            }
+        if (appTasks.find { getRunningAppState(it.key.id) == RunningAppState.RUNNING } != null) {
+            return RunningAppState.RUNNING
+        }
+        if (appTasks.find { getRunningAppState(it.key.id) == RunningAppState.MINIMIZED } != null) {
+            return RunningAppState.MINIMIZED
+        }
+        return RunningAppState.NOT_RUNNING
+    }
+
     /** Get the [RunningAppState] for the given task. */
     fun getRunningAppState(taskId: Int): RunningAppState {
         return when (taskId) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index 30e4e47..334ba6e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -297,9 +297,7 @@
                 Log.e(TAG, "Could not instantiate BubbleBarBubble for " + bubbleInfos.get(i));
                 continue;
             }
-            addBubbleInternally(bubble,  /* showAppBadge = */
-                    mBubbleBarViewController.isExpanded() || i == 0,
-                    /* isExpanding = */ false,  /* suppressAnimation = */ true);
+            addBubbleInternally(bubble, /* isExpanding= */ false, /* suppressAnimation= */ true);
         }
     }
 
@@ -390,8 +388,7 @@
             for (int i = update.currentBubbles.size() - 1; i >= 0; i--) {
                 BubbleBarBubble bubble = update.currentBubbles.get(i);
                 if (bubble != null) {
-                    addBubbleInternally(bubble, /* showAppBadge = */ !isCollapsed || i == 0,
-                            isExpanding, suppressAnimation);
+                    addBubbleInternally(bubble, isExpanding, suppressAnimation);
                     if (isCollapsed) {
                         // If we're collapsed, the most recently added bubble will be selected.
                         bubbleToSelect = bubble;
@@ -563,10 +560,8 @@
         }
     }
 
-    private void addBubbleInternally(BubbleBarBubble bubble, boolean showAppBadge,
-            boolean isExpanding, boolean suppressAnimation) {
-        //TODO(b/360652359): remove setting scale to the app badge once issue is fixed
-        bubble.getView().setBadgeScale(showAppBadge ? 1 : 0);
+    private void addBubbleInternally(BubbleBarBubble bubble, boolean isExpanding,
+            boolean suppressAnimation) {
         mBubbles.put(bubble.getKey(), bubble);
         mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index d91d10a..c0a76a8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -791,6 +791,7 @@
         updateLayoutParams();
         updateBubbleAccessibilityStates();
         updateContentDescription();
+        updateDotsAndBadgesIfCollapsed();
     }
 
     /** Removes the given bubble from the bubble bar. */
@@ -856,7 +857,7 @@
         updateBubbleAccessibilityStates();
         updateContentDescription();
         mDismissedByDragBubbleView = null;
-        updateNotificationDotsIfCollapsed();
+        updateDotsAndBadgesIfCollapsed();
     }
 
     /**
@@ -886,17 +887,23 @@
         return childViews;
     }
 
-    private void updateNotificationDotsIfCollapsed() {
+    private void updateDotsAndBadgesIfCollapsed() {
         if (isExpanded()) {
             return;
         }
         for (int i = 0; i < getChildCount(); i++) {
             BubbleView bubbleView = (BubbleView) getChildAt(i);
-            // when we're collapsed, the first bubble should show the dot if it has it. the rest of
-            // the bubbles should hide their dots.
-            if (i == 0 && bubbleView.hasUnseenContent()) {
-                bubbleView.showDotIfNeeded(/* animate= */ true);
+            // when we're collapsed, the first bubble should show the badge and the dot if it has
+            // it. the rest of the bubbles should hide their badges and dots.
+            if (i == 0) {
+                bubbleView.showBadge();
+                if (bubbleView.hasUnseenContent()) {
+                    bubbleView.showDotIfNeeded(/* animate= */ true);
+                } else {
+                    bubbleView.hideDot();
+                }
             } else {
+                bubbleView.hideBadge();
                 bubbleView.hideDot();
             }
         }
@@ -1100,7 +1107,7 @@
             }
             updateBubblesLayoutProperties(mBubbleBarLocation);
             updateContentDescription();
-            updateNotificationDotsIfCollapsed();
+            updateDotsAndBadgesIfCollapsed();
         }
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
index 114edf4..0ea4222 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
@@ -49,6 +49,8 @@
 public class BubbleView extends ConstraintLayout {
 
     public static final int DEFAULT_PATH_SIZE = 100;
+    /** Duration for animating the scale of the dot and badge. */
+    private static final int SCALE_ANIMATION_DURATION_MS = 200;
 
     private final ImageView mBubbleIcon;
     private final ImageView mAppIcon;
@@ -316,12 +318,37 @@
     }
 
     void setBadgeScale(float fraction) {
-        if (mAppIcon.getVisibility() == VISIBLE) {
+        if (hasBadge()) {
             mAppIcon.setScaleX(fraction);
             mAppIcon.setScaleY(fraction);
         }
     }
 
+    void showBadge() {
+        animateBadgeScale(1);
+    }
+
+    void hideBadge() {
+        animateBadgeScale(0);
+    }
+
+    private boolean hasBadge() {
+        return mAppIcon.getVisibility() == VISIBLE;
+    }
+
+    private void animateBadgeScale(float scale) {
+        if (!hasBadge()) {
+            return;
+        }
+        mAppIcon.clearAnimation();
+        mAppIcon.animate()
+                .setDuration(SCALE_ANIMATION_DURATION_MS)
+                .setInterpolator(Interpolators.FAST_OUT_SLOW_IN)
+                .scaleX(scale)
+                .scaleY(scale)
+                .start();
+    }
+
     /** Suppresses or un-suppresses drawing the dot due to an update for this bubble. */
     public void suppressDotForBubbleUpdate(boolean suppress) {
         mDotSuppressedForBubbleUpdate = suppress;
@@ -409,7 +436,7 @@
 
         clearAnimation();
         animate()
-                .setDuration(200)
+                .setDuration(SCALE_ANIMATION_DURATION_MS)
                 .setInterpolator(Interpolators.FAST_OUT_SLOW_IN)
                 .setUpdateListener((valueAnimator) -> {
                     float fraction = valueAnimator.getAnimatedFraction();
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 5b085d2..fcc5121 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -74,6 +74,7 @@
 import android.view.View;
 
 import androidx.annotation.BinderThread;
+import androidx.annotation.MainThread;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
@@ -550,11 +551,15 @@
                 @Override
                 public void onInputDeviceAdded(int deviceId) {
                     if (isTrackpadDevice(deviceId)) {
-                        boolean wasEmpty = mTrackpadsConnected.isEmpty();
-                        mTrackpadsConnected.add(deviceId);
-                        if (wasEmpty) {
-                            update();
-                        }
+                        // This updates internal TIS state so it needs to also run on the main
+                        // thread.
+                        MAIN_EXECUTOR.execute(() -> {
+                            boolean wasEmpty = mTrackpadsConnected.isEmpty();
+                            mTrackpadsConnected.add(deviceId);
+                            if (wasEmpty) {
+                                update();
+                            }
+                        });
                     }
                 }
 
@@ -564,12 +569,17 @@
 
                 @Override
                 public void onInputDeviceRemoved(int deviceId) {
-                    mTrackpadsConnected.remove(deviceId);
-                    if (mTrackpadsConnected.isEmpty()) {
-                        update();
-                    }
+                    // This updates internal TIS state so it needs to also run on the main
+                    // thread.
+                    MAIN_EXECUTOR.execute(() -> {
+                        mTrackpadsConnected.remove(deviceId);
+                        if (mTrackpadsConnected.isEmpty()) {
+                            update();
+                        }
+                    });
                 }
 
+                @MainThread
                 private void update() {
                     if (mInputMonitorCompat != null && !mTrackpadsConnected.isEmpty()) {
                         // Don't destroy and reinitialize input monitor due to trackpad
@@ -580,6 +590,7 @@
                 }
 
                 private boolean isTrackpadDevice(int deviceId) {
+                    // This is a blocking binder call that should run on a bg thread.
                     InputDevice inputDevice = mInputManager.getInputDevice(deviceId);
                     if (inputDevice == null) {
                         return false;
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
index 59413d3..066ddc0 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
@@ -32,6 +32,7 @@
 import com.android.launcher3.model.data.AppInfo
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.TaskItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.RecentsModel.RecentTasksChangedListener
 import com.android.quickstep.TaskIconCache
@@ -78,7 +79,9 @@
     private var taskListChangeId: Int = 1
 
     private lateinit var recentAppsController: TaskbarRecentAppsController
-    private lateinit var userHandle: UserHandle
+    private lateinit var myUserHandle: UserHandle
+    private val USER_HANDLE_1 = UserHandle.of(1)
+    private val USER_HANDLE_2 = UserHandle.of(2)
 
     private var canShowRunningAndRecentAppsAtInit = true
     private var recentTasksChangedListener: RecentTasksChangedListener? = null
@@ -86,7 +89,7 @@
     @Before
     fun setUp() {
         super.setup()
-        userHandle = Process.myUserHandle()
+        myUserHandle = Process.myUserHandle()
 
         // Set desktop mode supported
         whenever(mockContext.getResources()).thenReturn(mockResources)
@@ -149,6 +152,84 @@
     }
 
     @Test
+    fun getDesktopItemState_nullItemInfo_returnsNotRunning() {
+        setInDesktopMode(true)
+        assertThat(recentAppsController.getDesktopItemState(/* itemInfo= */ null))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_noItemPackage_returnsNotRunning() {
+        setInDesktopMode(true)
+        assertThat(recentAppsController.getDesktopItemState(ItemInfo()))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_noMatchingTasks_returnsNotRunning() {
+        setInDesktopMode(true)
+        val itemInfo = createItemInfo("package")
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_matchingVisibleTask_returnsVisible() {
+        setInDesktopMode(true)
+        val visibleTask = createTask(id = 1, "visiblePackage", isVisible = true)
+        updateRecentTasks(runningTasks = listOf(visibleTask), recentTaskPackages = emptyList())
+        val itemInfo = createItemInfo("visiblePackage")
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_matchingMinimizedTask_returnsMinimized() {
+        setInDesktopMode(true)
+        val minimizedTask = createTask(id = 1, "minimizedPackage", isVisible = false)
+        updateRecentTasks(runningTasks = listOf(minimizedTask), recentTaskPackages = emptyList())
+        val itemInfo = createItemInfo("minimizedPackage")
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.MINIMIZED)
+    }
+
+    @Test
+    fun getDesktopItemState_matchingMinimizedAndRunningTask_returnsVisible() {
+        setInDesktopMode(true)
+        updateRecentTasks(
+            runningTasks =
+                listOf(
+                    createTask(id = 1, "package", isVisible = false),
+                    createTask(id = 2, "package", isVisible = true),
+                ),
+            recentTaskPackages = emptyList(),
+        )
+        val itemInfo = createItemInfo("package")
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_noMatchingUserId_returnsNotRunning() {
+        setInDesktopMode(true)
+        updateRecentTasks(
+            runningTasks =
+                listOf(
+                    createTask(id = 1, "package", isVisible = false, USER_HANDLE_1),
+                    createTask(id = 2, "package", isVisible = true, USER_HANDLE_1),
+                ),
+            recentTaskPackages = emptyList(),
+        )
+        val itemInfo = createItemInfo("package", USER_HANDLE_2)
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
     fun getRunningAppState_taskNotRunningOrMinimized_returnsNotRunning() {
         setInDesktopMode(true)
         updateRecentTasks(runningTasks = emptyList(), recentTaskPackages = emptyList())
@@ -814,7 +895,13 @@
     private fun createTestAppInfo(
         packageName: String = "testPackageName",
         className: String = "testClassName",
-    ) = AppInfo(ComponentName(packageName, className), className /* title */, userHandle, Intent())
+    ) =
+        AppInfo(
+            ComponentName(packageName, className),
+            className /* title */,
+            myUserHandle,
+            Intent(),
+        )
 
     private fun createRecentTasksFromPackageNames(packageNames: List<String>): List<GroupTask> {
         return packageNames.map { packageName ->
@@ -833,14 +920,19 @@
         }
     }
 
-    private fun createTask(id: Int, packageName: String, isVisible: Boolean = true): Task {
+    private fun createTask(
+        id: Int,
+        packageName: String,
+        isVisible: Boolean = true,
+        localUserHandle: UserHandle? = null,
+    ): Task {
         return Task(
                 Task.TaskKey(
                     id,
                     WINDOWING_MODE_FREEFORM,
                     Intent().apply { `package` = packageName },
                     ComponentName(packageName, "TestActivity"),
-                    userHandle.identifier,
+                    localUserHandle?.identifier ?: myUserHandle.identifier,
                     0,
                 )
             )
@@ -852,6 +944,16 @@
             .thenReturn(inDesktopMode)
     }
 
+    private fun createItemInfo(
+        packageName: String,
+        userHandle: UserHandle = myUserHandle,
+    ): ItemInfo {
+        val appInfo = AppInfo()
+        appInfo.intent = Intent().setComponent(ComponentName(packageName, "className"))
+        appInfo.user = userHandle
+        return WorkspaceItemInfo(appInfo)
+    }
+
     private val GroupTask.packageNames: List<String>
         get() = tasks.map { task -> task.key.packageName }
 
diff --git a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
index 44c23ba..6a7b6f8 100644
--- a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
+++ b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
@@ -26,6 +26,7 @@
 import com.android.launcher3.tapl.LaunchedAppState;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
+import com.android.launcher3.util.TestUtil;
 import com.android.quickstep.views.RecentsView;
 
 import org.junit.rules.RuleChain;
@@ -56,7 +57,7 @@
     protected void assertTestActivityIsRunning(int activityNumber, String message) {
         assertTrue(message, mDevice.wait(
                 Until.hasObject(By.pkg(getAppPackageName()).text("TestActivity" + activityNumber)),
-                DEFAULT_UI_TIMEOUT));
+                TestUtil.DEFAULT_UI_TIMEOUT));
     }
 
     protected LaunchedAppState getAndAssertLaunchedApp() {
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 1f11c14..aa105f9 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -22,9 +22,7 @@
 import static com.android.launcher3.tapl.LauncherInstrumentation.WAIT_TIME_MS;
 import static com.android.launcher3.tapl.TestHelpers.getHomeIntentInPackage;
 import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
-import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_ACTIVITY_TIMEOUT;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_BROADCAST_TIMEOUT_SECS;
-import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_UI_TIMEOUT;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.resolveSystemApp;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.startAppFast;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.startTestActivity;
@@ -56,6 +54,7 @@
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ExtendedLongPressTimeoutRule;
 import com.android.launcher3.util.rule.FailureWatcher;
@@ -214,7 +213,7 @@
             }
             result[0] = f.apply(activity);
             return true;
-        }).get(), DEFAULT_UI_TIMEOUT, mLauncher);
+        }).get(), mLauncher);
         return (T) result[0];
     }
 
@@ -244,7 +243,7 @@
 
         Wait.atMost("Recents activity didn't stop",
                 () -> getFromRecents(recents -> !recents.isStarted()),
-                DEFAULT_UI_TIMEOUT, mLauncher);
+                mLauncher);
     }
 
     @Test
@@ -254,7 +253,8 @@
         startTestActivity(2);
         waitForRecentsActivityStop();
         Wait.atMost("Expected three apps in the task list",
-                () -> mLauncher.getRecentTasks().size() >= 3, DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+                () -> mLauncher.getRecentTasks().size() >= 3,
+                mLauncher);
 
         checkTestLauncher();
         BaseOverview overview = mLauncher.getLaunchedAppState().switchToOverview();
@@ -282,7 +282,7 @@
         assertNotNull("OverviewTask.open returned null", task.open());
         assertTrue("Test activity didn't open from Overview", TestHelpers.wait(Until.hasObject(
                 By.pkg(getAppPackageName()).text("TestActivity2")),
-                DEFAULT_UI_TIMEOUT));
+                TestUtil.DEFAULT_UI_TIMEOUT));
 
 
         // Test dismissing a task.
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 4459ed6..77f4c05 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -57,8 +57,6 @@
 
     static final String TAG = "QuickStepOnOffRule";
 
-    public static final int WAIT_TIME_MS = 10000;
-
     public enum Mode {
         THREE_BUTTON, ZERO_BUTTON, ALL
     }
@@ -179,12 +177,13 @@
         }
 
         Wait.atMost("Couldn't switch to " + overlayPackage,
-                () -> launcher.getNavigationModel() == expectedMode, WAIT_TIME_MS, launcher);
+                () -> launcher.getNavigationModel() == expectedMode,
+                launcher);
 
         Wait.atMost(() -> "Switching nav mode: "
                         + launcher.getNavigationModeMismatchError(false),
                 () -> launcher.getNavigationModeMismatchError(false) == null,
-                WAIT_TIME_MS, launcher);
+                launcher);
         AbstractLauncherUiTest.checkDetectedLeaks(launcher, false);
         return true;
     }
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt b/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt
index 120a89b..f58c84e 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt
@@ -26,6 +26,7 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape
 import com.android.launcher3.uioverrides.QuickstepLauncher
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertWithMessage
 import org.junit.Before
 import org.junit.Test
@@ -172,7 +173,7 @@
             .that(
                 mDevice.wait(
                     Until.hasObject(By.pkg(getAppPackageName()).text("TestActivity$index")),
-                    DEFAULT_UI_TIMEOUT,
+                    TestUtil.DEFAULT_UI_TIMEOUT,
                 )
             )
             .isTrue()
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 5ff2af7..f1fe2d2 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -47,6 +47,7 @@
 import com.android.launcher3.tapl.SelectModeButtons;
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.TestStabilityRule;
@@ -145,7 +146,7 @@
         assertNotNull("OverviewTask.open returned null", task.open());
         assertTrue("Test activity didn't open from Overview", mDevice.wait(Until.hasObject(
                         By.pkg(getAppPackageName()).text("TestActivity2")),
-                DEFAULT_UI_TIMEOUT));
+                TestUtil.DEFAULT_UI_TIMEOUT));
         executeOnLauncher(launcher -> assertTrue(
                 "Launcher activity is the top activity; expecting another activity to be the top "
                         + "one",
@@ -448,7 +449,7 @@
                 mDevice.wait(Until.hasObject(By.pkg(getAppPackageName()).text(
                                 mLauncher.isGridOnlyOverviewEnabled() ? "TestActivity12"
                                         : "TestActivity13")),
-                        DEFAULT_UI_TIMEOUT));
+                        TestUtil.DEFAULT_UI_TIMEOUT));
 
         // Scroll the task offscreen as it is now first
         overview = mLauncher.goHome().switchToOverview();
@@ -563,7 +564,7 @@
             mLauncher.getDevice().setOrientationLeft();
             startTestActivity(7);
             Wait.atMost("Device should not be in natural orientation",
-                    () -> !mDevice.isNaturalOrientation(), DEFAULT_UI_TIMEOUT, mLauncher);
+                    () -> !mDevice.isNaturalOrientation(), mLauncher);
             mLauncher.goHome();
         } finally {
             mLauncher.setExpectedRotationCheckEnabled(true);
diff --git a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
index fc8465d..1b58987 100644
--- a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
@@ -95,6 +95,7 @@
 import com.android.launcher3.views.ScrimView;
 import com.android.launcher3.views.SpringRelativeLayout;
 import com.android.launcher3.workprofile.PersonalWorkSlidingTabStrip;
+import com.android.systemui.plugins.AllAppsRow;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -152,6 +153,7 @@
     private final RectF mTmpRectF = new RectF();
     protected AllAppsPagedView mViewPager;
     protected FloatingHeaderView mHeader;
+    protected final List<AllAppsRow> mAdditionalHeaderRows = new ArrayList<>();
     protected View mBottomSheetBackground;
     protected RecyclerViewFastScroller mFastScroller;
     private ConstraintLayout mFastScrollLetterLayout;
@@ -262,6 +264,8 @@
 
         getLayoutInflater().inflate(R.layout.all_apps_content, this);
         mHeader = findViewById(R.id.all_apps_header);
+        mAdditionalHeaderRows.clear();
+        mAdditionalHeaderRows.addAll(getAdditionalHeaderRows());
         mBottomSheetBackground = findViewById(R.id.bottom_sheet_background);
         mBottomSheetHandleArea = findViewById(R.id.bottom_sheet_handle_area);
         mSearchRecyclerView = findViewById(R.id.search_results_list_view);
@@ -288,6 +292,10 @@
         mSearchUiManager = (SearchUiManager) mSearchContainer;
     }
 
+    public List<AllAppsRow> getAdditionalHeaderRows() {
+        return List.of();
+    }
+
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
@@ -714,6 +722,8 @@
     }
 
     void setupHeader() {
+        mAdditionalHeaderRows.forEach(row -> mHeader.onPluginDisconnected(row));
+
         mHeader.setVisibility(View.VISIBLE);
         boolean tabsHidden = !mUsingTabs;
         mHeader.setup(
@@ -731,6 +741,7 @@
                 adapterHolder.mRecyclerView.scrollToTop();
             }
         });
+        mAdditionalHeaderRows.forEach(row -> mHeader.onPluginConnected(row, mActivityContext));
 
         removeCustomRules(mHeader);
         if (isSearchBarFloating()) {
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index ac06ab4..8193511 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.allapps;
 
+import static com.android.launcher3.allapps.FloatingHeaderRow.NO_ROWS;
+
 import android.animation.ValueAnimator;
 import android.content.Context;
 import android.graphics.Point;
@@ -109,11 +111,11 @@
 
     // This is initialized once during inflation and stays constant after that. Fixed views
     // cannot be added or removed dynamically.
-    private FloatingHeaderRow[] mFixedRows = FloatingHeaderRow.NO_ROWS;
+    private FloatingHeaderRow[] mFixedRows = NO_ROWS;
 
     // Array of all fixed rows and plugin rows. This is initialized every time a plugin is
     // enabled or disabled, and represent the current set of all rows.
-    private FloatingHeaderRow[] mAllRows = FloatingHeaderRow.NO_ROWS;
+    private FloatingHeaderRow[] mAllRows = NO_ROWS;
 
     public FloatingHeaderView(@NonNull Context context) {
         this(context, null);
@@ -180,6 +182,10 @@
 
     @Override
     public void onPluginConnected(AllAppsRow allAppsRowPlugin, Context context) {
+        if (mPluginRows.containsKey(allAppsRowPlugin)) {
+            // Plugin has already been connected
+            return;
+        }
         PluginHeaderRow headerRow = new PluginHeaderRow(allAppsRowPlugin, this);
         addView(headerRow.mView, indexOfChild(mTabLayout));
         mPluginRows.put(allAppsRowPlugin, headerRow);
@@ -211,6 +217,9 @@
     @Override
     public void onPluginDisconnected(AllAppsRow plugin) {
         PluginHeaderRow row = mPluginRows.get(plugin);
+        if (row == null) {
+            return;
+        }
         removeView(row.mView);
         mPluginRows.remove(plugin);
         recreateAllRowsArray();
diff --git a/tests/src/com/android/launcher3/LauncherIntentTest.java b/tests/src/com/android/launcher3/LauncherIntentTest.java
index aeeb42a..3e16713 100644
--- a/tests/src/com/android/launcher3/LauncherIntentTest.java
+++ b/tests/src/com/android/launcher3/LauncherIntentTest.java
@@ -23,7 +23,7 @@
 import android.platform.test.annotations.LargeTest;
 import android.view.KeyEvent;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.launcher3.allapps.ActivityAllAppsContainerView;
 import com.android.launcher3.allapps.SearchRecyclerView;
diff --git a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
index 44b8ff8..1816030 100644
--- a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
@@ -69,8 +69,7 @@
     private void verifyAppUninstalledFromAllApps(Workspace workspace, String appName) {
         final HomeAllApps allApps = workspace.switchToAllApps();
         Wait.atMost(appName + " app was found on all apps after being uninstalled",
-                () -> allApps.tryGetAppIcon(appName) == null,
-                DEFAULT_UI_TIMEOUT, mLauncher);
+                () -> allApps.tryGetAppIcon(appName) == null, mLauncher);
     }
 
     private void installDummyAppAndWaitForUIUpdate() throws IOException {
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 206647a..1fbdceb 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -32,8 +32,6 @@
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
-import android.content.pm.LauncherActivityInfo;
-import android.content.pm.LauncherApps;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.graphics.Point;
@@ -97,10 +95,8 @@
  */
 public abstract class AbstractLauncherUiTest<LAUNCHER_TYPE extends Launcher> {
 
-    public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 10;
 
-    public static final long DEFAULT_UI_TIMEOUT = TestUtil.DEFAULT_UI_TIMEOUT;
     private static final String TAG = "AbstractLauncherUiTest";
 
     private static final long BYTES_PER_MEGABYTE = 1 << 20;
@@ -151,7 +147,7 @@
                     launcher.forceGc();
                     return MAIN_EXECUTOR.submit(
                             () -> launcher.noLeakedActivities(requireOneActiveActivity)).get();
-                }, DEFAULT_UI_TIMEOUT, launcher);
+                }, launcher);
     }
 
     public static String getAppPackageName() {
@@ -443,7 +439,7 @@
      */
     protected <T> T getOnUiThread(final Callable<T> callback) {
         try {
-            return mMainThreadExecutor.submit(callback).get(DEFAULT_UI_TIMEOUT,
+            return mMainThreadExecutor.submit(callback).get(TestUtil.DEFAULT_UI_TIMEOUT,
                     TimeUnit.MILLISECONDS);
         } catch (TimeoutException e) {
             Log.e(TAG, "Timeout in getOnUiThread, sending SIGABRT", e);
@@ -498,13 +494,7 @@
     // flakiness.
     protected void waitForLauncherCondition(String
             message, Function<LAUNCHER_TYPE, Boolean> condition) {
-        waitForLauncherCondition(message, condition, DEFAULT_ACTIVITY_TIMEOUT);
-    }
-
-    // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
-    // flakiness.
-    protected <O> O getOnceNotNull(String message, Function<LAUNCHER_TYPE, O> f) {
-        return getOnceNotNull(message, f, DEFAULT_ACTIVITY_TIMEOUT);
+        waitForLauncherCondition(message, condition, TestUtil.DEFAULT_UI_TIMEOUT);
     }
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
@@ -513,12 +503,12 @@
             String message, Function<LAUNCHER_TYPE, Boolean> condition, long timeout) {
         verifyKeyguardInvisible();
         if (!TestHelpers.isInLauncherProcess()) return;
-        Wait.atMost(message, () -> getFromLauncher(condition), timeout, mLauncher);
+        Wait.atMost(message, () -> getFromLauncher(condition), mLauncher, timeout);
     }
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
-    protected <T> T getOnceNotNull(String message, Function<LAUNCHER_TYPE, T> f, long timeout) {
+    protected <T> T getOnceNotNull(String message, Function<LAUNCHER_TYPE, T> f) {
         if (!TestHelpers.isInLauncherProcess()) return null;
 
         final Object[] output = new Object[1];
@@ -526,7 +516,7 @@
             final Object fromLauncher = getFromLauncher(f);
             output[0] = fromLauncher;
             return fromLauncher != null;
-        }, timeout, mLauncher);
+        }, mLauncher);
         return (T) output[0];
     }
 
@@ -540,12 +530,7 @@
         Wait.atMost(message, () -> {
             testThreadAction.run();
             return getFromLauncher(condition);
-        }, timeout, mLauncher);
-    }
-
-    protected LauncherActivityInfo getSettingsApp() {
-        return mTargetContext.getSystemService(LauncherApps.class)
-                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
+        }, mLauncher, timeout);
     }
 
     /**
@@ -633,13 +618,13 @@
         }
         getInstrumentation().getTargetContext().startActivity(intent);
         assertTrue("App didn't start: " + selector,
-                TestHelpers.wait(Until.hasObject(selector), DEFAULT_UI_TIMEOUT));
+                TestHelpers.wait(Until.hasObject(selector), TestUtil.DEFAULT_UI_TIMEOUT));
 
         // Wait for the Launcher to stop.
         final LauncherInstrumentation launcherInstrumentation = new LauncherInstrumentation();
         Wait.atMost("Launcher activity didn't stop",
                 () -> !launcherInstrumentation.isLauncherActivityStarted(),
-                DEFAULT_ACTIVITY_TIMEOUT, launcherInstrumentation);
+                launcherInstrumentation);
     }
 
     public static ActivityInfo resolveSystemAppInfo(String category) {
@@ -662,8 +647,7 @@
                 launcher.finish();
             }
         });
-        waitForLauncherCondition(
-                "Launcher still active", launcher -> launcher == null, DEFAULT_UI_TIMEOUT);
+        waitForLauncherCondition("Launcher still active", launcher -> launcher == null);
     }
 
     protected boolean isInLaunchedApp(LAUNCHER_TYPE launcher) {
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
index 7ff4f22..7845222 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
@@ -103,12 +103,12 @@
 
         setResultAndWaitForAnimation(acceptConfig);
         if (acceptConfig) {
-            Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+            Wait.atMost("", new WidgetSearchCondition(), mLauncher);
             assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
         } else {
             // Verify that the widget id is deleted.
             Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
-                    DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+                    mLauncher);
         }
     }
 
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
index 9a2147a..460ffc4 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
@@ -29,6 +29,7 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
 import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
@@ -68,7 +69,7 @@
         resizeFrame.dismiss();
 
         final Widget widget = mLauncher.getWorkspace().tryGetWidget(widgetInfo.label,
-                DEFAULT_UI_TIMEOUT);
+                TestUtil.DEFAULT_UI_TIMEOUT);
         assertNotNull("Widget not found on the workspace", widget);
         widget.launch(getAppPackageName());
         mLauncher.disableDebugTracing(); // b/289161193
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
index d40d3bc..4cdbd96 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
@@ -53,6 +53,7 @@
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
@@ -233,13 +234,15 @@
     }
 
     private void verifyWidgetPresent(LauncherAppWidgetProviderInfo info) {
-        final Widget widget = mLauncher.getWorkspace().tryGetWidget(info.label, DEFAULT_UI_TIMEOUT);
+        final Widget widget = mLauncher.getWorkspace().tryGetWidget(info.label,
+                TestUtil.DEFAULT_UI_TIMEOUT);
         assertTrue("Widget is not present",
                 widget != null);
     }
 
     private void verifyPendingWidgetPresent() {
-        final Widget widget = mLauncher.getWorkspace().tryGetPendingWidget(DEFAULT_UI_TIMEOUT);
+        final Widget widget = mLauncher.getWorkspace().tryGetPendingWidget(
+                TestUtil.DEFAULT_UI_TIMEOUT);
         assertTrue("Pending widget is not present",
                 widget != null);
     }
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
index 35c7cab..fe3b2ee 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
@@ -169,8 +169,7 @@
 
         // Go back to home
         mLauncher.goHome();
-        Wait.atMost("", new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT,
-                mLauncher);
+        Wait.atMost("", new ItemSearchCondition(itemMatcher), mLauncher);
     }
 
     /**
diff --git a/tests/src/com/android/launcher3/util/Wait.java b/tests/src/com/android/launcher3/util/Wait.java
deleted file mode 100644
index 50bc32e..0000000
--- a/tests/src/com/android/launcher3/util/Wait.java
+++ /dev/null
@@ -1,66 +0,0 @@
-package com.android.launcher3.util;
-
-import android.os.SystemClock;
-import android.util.Log;
-
-import com.android.launcher3.tapl.LauncherInstrumentation;
-
-import org.junit.Assert;
-
-import java.util.function.Supplier;
-
-/**
- * A utility class for waiting for a condition to be true.
- */
-public class Wait {
-
-    private static final long DEFAULT_SLEEP_MS = 200;
-
-    public static void atMost(String message, Condition condition, long timeout,
-            LauncherInstrumentation launcher) {
-        atMost(() -> message, condition, timeout, DEFAULT_SLEEP_MS, launcher);
-    }
-
-    public static void atMost(Supplier<String> message, Condition condition, long timeout,
-            LauncherInstrumentation launcher) {
-        atMost(message, condition, timeout, DEFAULT_SLEEP_MS, launcher);
-    }
-
-    public static void atMost(Supplier<String> message, Condition condition, long timeout,
-            long sleepMillis,
-            LauncherInstrumentation launcher) {
-        final long startTime = SystemClock.uptimeMillis();
-        long endTime = startTime + timeout;
-        Log.d("Wait", "atMost: " + startTime + " - " + endTime);
-        while (SystemClock.uptimeMillis() < endTime) {
-            try {
-                if (condition.isTrue()) {
-                    return;
-                }
-            } catch (Throwable t) {
-                throw new RuntimeException(t);
-            }
-            SystemClock.sleep(sleepMillis);
-        }
-
-        // Check once more before returning false.
-        try {
-            if (condition.isTrue()) {
-                return;
-            }
-        } catch (Throwable t) {
-            throw new RuntimeException(t);
-        }
-        Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis());
-        launcher.checkForAnomaly(false, false);
-        Assert.fail(message.get());
-    }
-
-    /**
-     * Interface representing a generic condition
-     */
-    public interface Condition {
-
-        boolean isTrue() throws Throwable;
-    }
-}
diff --git a/tests/src/com/android/launcher3/util/Wait.kt b/tests/src/com/android/launcher3/util/Wait.kt
new file mode 100644
index 0000000..1e5af54
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/Wait.kt
@@ -0,0 +1,79 @@
+/*
+ * 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.launcher3.util
+
+import android.os.SystemClock
+import android.util.Log
+import com.android.launcher3.tapl.LauncherInstrumentation
+import java.util.function.Supplier
+import org.junit.Assert
+
+/** A utility class for waiting for a condition to be true. */
+object Wait {
+    private const val DEFAULT_SLEEP_MS: Long = 200
+
+    @JvmStatic
+    @JvmOverloads
+    fun atMost(
+        message: String,
+        condition: Condition,
+        launcherInstrumentation: LauncherInstrumentation? = null,
+        timeout: Long = TestUtil.DEFAULT_UI_TIMEOUT,
+    ) {
+        atMost({ message }, condition, launcherInstrumentation, timeout)
+    }
+
+    @JvmStatic
+    @JvmOverloads
+    fun atMost(
+        message: Supplier<String>,
+        condition: Condition,
+        launcherInstrumentation: LauncherInstrumentation? = null,
+        timeout: Long = TestUtil.DEFAULT_UI_TIMEOUT,
+    ) {
+        val startTime = SystemClock.uptimeMillis()
+        val endTime = startTime + timeout
+        Log.d("Wait", "atMost: $startTime - $endTime")
+        while (SystemClock.uptimeMillis() < endTime) {
+            try {
+                if (condition.isTrue()) {
+                    return
+                }
+            } catch (t: Throwable) {
+                throw RuntimeException(t)
+            }
+            SystemClock.sleep(DEFAULT_SLEEP_MS)
+        }
+
+        // Check once more before returning false.
+        try {
+            if (condition.isTrue()) {
+                return
+            }
+        } catch (t: Throwable) {
+            throw RuntimeException(t)
+        }
+        Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis())
+        launcherInstrumentation?.checkForAnomaly(false, false)
+        Assert.fail(message.get())
+    }
+
+    /** Interface representing a generic condition */
+    fun interface Condition {
+
+        @Throws(Throwable::class) fun isTrue(): Boolean
+    }
+}