diff --git a/Android.bp b/Android.bp
index d04dca0..f79c186 100644
--- a/Android.bp
+++ b/Android.bp
@@ -258,8 +258,8 @@
         "go/quickstep/res",
     ],
     static_libs: [
-        "Launcher3CommonDepsLib",
         "QuickstepResLib",
+        "Launcher3CommonDepsLib",
     ],
     manifest: "quickstep/AndroidManifest-launcher.xml",
     additional_manifests: [
@@ -278,16 +278,15 @@
     srcs: [
         ":launcher-src-no-build-config",
     ],
-    resource_dirs: [
-        "quickstep/res",
-    ],
+    resource_dirs: [],
     libs: [
         "framework-statsd",
     ],
     static_libs: [
+        "QuickstepResLib",
         "SystemUI-statsd",
         "SystemUISharedLib",
-        "Launcher3CommonDepsLib"
+        "Launcher3CommonDepsLib",
     ],
     manifest: "quickstep/AndroidManifest.xml",
     platform_apis: true,
diff --git a/quickstep/res/drawable/task_menu_item_bg.xml b/quickstep/res/drawable/task_menu_item_bg.xml
index b6a8b90..16c13eb 100644
--- a/quickstep/res/drawable/task_menu_item_bg.xml
+++ b/quickstep/res/drawable/task_menu_item_bg.xml
@@ -15,7 +15,8 @@
      limitations under the License.
 -->
 
-<shape xmlns:android="http://schemas.android.com/apk/res/android">
-    <solid android:color="?android:attr/colorPrimary"/>
-    <corners android:radius="@dimen/task_menu_item_corner_radius"/>
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:androidprv="http://schemas.android.com/apk/prv/res/android">
+    <solid android:color="?androidprv:attr/colorSurface" />
+    <corners android:radius="@dimen/task_menu_item_corner_radius" />
 </shape>
diff --git a/quickstep/res/layout/task_menu_with_arrow.xml b/quickstep/res/layout/task_menu_with_arrow.xml
new file mode 100644
index 0000000..38573fd
--- /dev/null
+++ b/quickstep/res/layout/task_menu_with_arrow.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 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.
+-->
+<com.android.quickstep.views.TaskMenuViewWithArrow
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:animateLayoutChanges="true"
+    android:background="@drawable/task_menu_bg"
+    android:orientation="vertical"
+    android:visibility="invisible">
+
+    <LinearLayout
+        android:id="@+id/menu_option_layout"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:orientation="vertical"
+        android:showDividers="middle" />
+
+</com.android.quickstep.views.TaskMenuViewWithArrow>
\ No newline at end of file
diff --git a/quickstep/res/layout/task_view_menu_option.xml b/quickstep/res/layout/task_view_menu_option.xml
index 5978b97..8a8fc36 100644
--- a/quickstep/res/layout/task_view_menu_option.xml
+++ b/quickstep/res/layout/task_view_menu_option.xml
@@ -18,7 +18,7 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"
-    android:orientation="vertical"
+    android:orientation="horizontal"
     android:paddingTop="@dimen/task_card_menu_option_vertical_padding"
     android:paddingBottom="@dimen/task_card_menu_option_vertical_padding"
     android:background="@drawable/task_menu_item_bg"
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 98d43f1..8649a1d 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -27,7 +27,7 @@
     <dimen name="task_menu_corner_radius">22dp</dimen>
     <dimen name="task_menu_item_corner_radius">4dp</dimen>
     <dimen name="task_menu_spacing">2dp</dimen>
-    <dimen name="task_menu_width_grid">200dp</dimen>
+    <dimen name="task_menu_width_grid">234dp</dimen>
     <dimen name="overview_proactive_row_height">48dp</dimen>
     <dimen name="overview_proactive_row_bottom_margin">16dp</dimen>
 
@@ -90,7 +90,7 @@
     <dimen name="task_menu_vertical_padding">8dp</dimen>
     <dimen name="task_card_margin">8dp</dimen>
     <dimen name="task_card_menu_shadow_height">3dp</dimen>
-    <dimen name="task_menu_option_start_margin">12dp</dimen>
+    <dimen name="task_menu_option_start_margin">16dp</dimen>
     <!-- Copied from framework resource:
        docked_stack_divider_thickness - 2 * docked_stack_divider_insets -->
     <dimen name="multi_window_task_divider_size">10dp</dimen>
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
index 6fbef9b..a8a0b59 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
@@ -102,9 +102,7 @@
     }
 
     private void showIMESwitcher() {
-        mService.getSystemService(InputMethodManager.class)
-                .showInputMethodPickerFromSystem(true /* showAuxiliarySubtypes */,
-                        DEFAULT_DISPLAY);
+        SystemUiProxy.INSTANCE.getNoCreate().onImeSwitcherPressed();
     }
 
     private void notifyImeClick(boolean longClick) {
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 3080f04..79dd633 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -22,35 +22,33 @@
 import android.app.ActivityManager;
 import android.os.Build;
 import android.os.Process;
-import android.util.Log;
+import android.os.RemoteException;
 import android.util.SparseBooleanArray;
 
 import androidx.annotation.VisibleForTesting;
 
-import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.util.LooperExecutor;
+import com.android.systemui.shared.recents.model.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.KeyguardManagerCompat;
-import com.android.systemui.shared.system.TaskStackChangeListener;
-import com.android.systemui.shared.system.TaskStackChangeListeners;
+import com.android.wm.shell.recents.IRecentTasksListener;
+import com.android.wm.shell.util.GroupedRecentTaskInfo;
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.function.Consumer;
 
 /**
  * Manages the recent task list from the system, caching it as necessary.
  */
 @TargetApi(Build.VERSION_CODES.R)
-public class RecentTasksList extends TaskStackChangeListener {
+public class RecentTasksList {
 
     private static final TaskLoadResult INVALID_RESULT = new TaskLoadResult(-1, false, 0);
 
     private final KeyguardManagerCompat mKeyguardManager;
     private final LooperExecutor mMainThreadExecutor;
-    private final ActivityManagerWrapper mActivityManagerWrapper;
+    private final SystemUiProxy mSysUiProxy;
 
     // The list change id, increments as the task list changes in the system
     private int mChangeId;
@@ -62,12 +60,17 @@
     private TaskLoadResult mResultsUi = INVALID_RESULT;
 
     public RecentTasksList(LooperExecutor mainThreadExecutor,
-            KeyguardManagerCompat keyguardManager, ActivityManagerWrapper activityManagerWrapper) {
+            KeyguardManagerCompat keyguardManager, SystemUiProxy sysUiProxy) {
         mMainThreadExecutor = mainThreadExecutor;
         mKeyguardManager = keyguardManager;
         mChangeId = 1;
-        mActivityManagerWrapper = activityManagerWrapper;
-        TaskStackChangeListeners.getInstance().registerTaskStackListener(this);
+        mSysUiProxy = sysUiProxy;
+        sysUiProxy.registerRecentTasksListener(new IRecentTasksListener.Stub() {
+            @Override
+            public void onRecentTasksChanged() throws RemoteException {
+                mMainThreadExecutor.execute(RecentTasksList.this::onRecentTasksChanged);
+            }
+        });
     }
 
     @VisibleForTesting
@@ -78,10 +81,11 @@
     /**
      * Fetches the task keys skipping any local cache.
      */
-    public void getTaskKeys(int numTasks, Consumer<ArrayList<Task>> callback) {
+    public void getTaskKeys(int numTasks, Consumer<ArrayList<GroupTask>> callback) {
         // Kick off task loading in the background
         UI_HELPER_EXECUTOR.execute(() -> {
-            ArrayList<Task> tasks = loadTasksInBackground(numTasks, -1, true /* loadKeysOnly */);
+            ArrayList<GroupTask> tasks = loadTasksInBackground(numTasks, -1,
+                    true /* loadKeysOnly */);
             mMainThreadExecutor.execute(() -> callback.accept(tasks));
         });
     }
@@ -93,14 +97,15 @@
      * @param callback The callback to receive the list of recent tasks
      * @return The change id of the current task list
      */
-    public synchronized int getTasks(boolean loadKeysOnly, Consumer<ArrayList<Task>> callback) {
+    public synchronized int getTasks(boolean loadKeysOnly,
+            Consumer<ArrayList<GroupTask>> callback) {
         final int requestLoadId = mChangeId;
         if (mResultsUi.isValidForRequest(requestLoadId, loadKeysOnly)) {
             // The list is up to date, send the callback on the next frame,
             // so that requestID can be returned first.
             if (callback != null) {
                 // Copy synchronously as the changeId might change by next frame
-                ArrayList<Task> result = copyOf(mResultsUi);
+                ArrayList<GroupTask> result = copyOf(mResultsUi);
                 mMainThreadExecutor.post(() -> {
                     callback.accept(result);
                 });
@@ -120,7 +125,7 @@
                 mLoadingTasksInBackground = false;
                 mResultsUi = loadResult;
                 if (callback != null) {
-                    ArrayList<Task> result = copyOf(mResultsUi);
+                    ArrayList<GroupTask> result = copyOf(mResultsUi);
                     callback.accept(result);
                 }
             });
@@ -136,35 +141,7 @@
         return mChangeId == changeId;
     }
 
-    @Override
-    public void onTaskStackChanged() {
-        invalidateLoadedTasks();
-    }
-
-    @Override
-    public void onRecentTaskListUpdated() {
-        // In some cases immediately after booting, the tasks in the system recent task list may be
-        // loaded, but not in the active task hierarchy in the system.  These tasks are displayed in 
-        // overview, but removing them don't result in a onTaskStackChanged() nor a onTaskRemoved()
-        // callback (those are for changes to the active tasks), but the task list is still updated,
-        // so we should also invalidate the change id to ensure we load a new list instead of 
-        // reusing a stale list.
-        invalidateLoadedTasks();
-    }
-
-    @Override
-    public void onTaskRemoved(int taskId) {
-        invalidateLoadedTasks();
-    }
-
-
-    @Override
-    public void onActivityPinned(String packageName, int userId, int taskId, int stackId) {
-        invalidateLoadedTasks();
-    }
-
-    @Override
-    public synchronized void onActivityUnpinned() {
+    public void onRecentTasksChanged() {
         invalidateLoadedTasks();
     }
 
@@ -180,8 +157,8 @@
     @VisibleForTesting
     TaskLoadResult loadTasksInBackground(int numTasks, int requestId, boolean loadKeysOnly) {
         int currentUserId = Process.myUserHandle().getIdentifier();
-        List<ActivityManager.RecentTaskInfo> rawTasks =
-                mActivityManagerWrapper.getRecentTasks(numTasks, currentUserId);
+        ArrayList<GroupedRecentTaskInfo> rawTasks =
+                mSysUiProxy.getRecentTasks(numTasks, currentUserId);
         // The raw tasks are given in most-recent to least-recent order, we need to reverse it
         Collections.reverse(rawTasks);
 
@@ -197,45 +174,53 @@
         };
 
         TaskLoadResult allTasks = new TaskLoadResult(requestId, loadKeysOnly, rawTasks.size());
-        for (ActivityManager.RecentTaskInfo rawTask : rawTasks) {
-            Task.TaskKey taskKey = new Task.TaskKey(rawTask);
-            Task task;
-            if (!loadKeysOnly) {
-                boolean isLocked = tmpLockedUsers.get(taskKey.userId);
-                task = Task.from(taskKey, rawTask, isLocked);
-            } else {
-                task = new Task(taskKey);
+        for (GroupedRecentTaskInfo rawTask : rawTasks) {
+            ActivityManager.RecentTaskInfo taskInfo1 = rawTask.mTaskInfo1;
+            ActivityManager.RecentTaskInfo taskInfo2 = rawTask.mTaskInfo2;
+            Task.TaskKey task1Key = new Task.TaskKey(taskInfo1);
+            Task task1 = loadKeysOnly
+                    ? new Task(task1Key)
+                    : Task.from(task1Key, taskInfo1,
+                            tmpLockedUsers.get(task1Key.userId) /* isLocked */);
+            task1.setLastSnapshotData(taskInfo1);
+            Task task2 = null;
+            if (taskInfo2 != null) {
+                Task.TaskKey task2Key = new Task.TaskKey(taskInfo2);
+                task2 = loadKeysOnly
+                        ? new Task(task2Key)
+                        : Task.from(task2Key, taskInfo2,
+                                tmpLockedUsers.get(task2Key.userId) /* isLocked */);
+                task2.setLastSnapshotData(taskInfo2);
             }
-            task.setLastSnapshotData(rawTask);
-            allTasks.add(task);
+            allTasks.add(new GroupTask(task1, task2));
         }
 
         return allTasks;
     }
 
-    private ArrayList<Task> copyOf(ArrayList<Task> tasks) {
-        ArrayList<Task> newTasks = new ArrayList<>();
+    private ArrayList<GroupTask> copyOf(ArrayList<GroupTask> tasks) {
+        ArrayList<GroupTask> newTasks = new ArrayList<>();
         for (int i = 0; i < tasks.size(); i++) {
-            newTasks.add(new Task(tasks.get(i)));
+            newTasks.add(new GroupTask(tasks.get(i)));
         }
         return newTasks;
     }
 
-    private static class TaskLoadResult extends ArrayList<Task> {
+    private static class TaskLoadResult extends ArrayList<GroupTask> {
 
-        final int mId;
+        final int mRequestId;
 
         // If the result was loaded with keysOnly  = true
         final boolean mKeysOnly;
 
-        TaskLoadResult(int id, boolean keysOnly, int size) {
+        TaskLoadResult(int requestId, boolean keysOnly, int size) {
             super(size);
-            mId = id;
+            mRequestId = requestId;
             mKeysOnly = keysOnly;
         }
 
         boolean isValidForRequest(int requestId, boolean loadKeysOnly) {
-            return mId == requestId && (!mKeysOnly || loadKeysOnly);
+            return mRequestId == requestId && (!mKeysOnly || loadKeysOnly);
         }
     }
 }
\ No newline at end of file
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index 1e82c8c..71153c6 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -34,6 +34,7 @@
 import com.android.launcher3.icons.IconProvider.IconChangeListener;
 import com.android.launcher3.util.Executors.SimpleThreadFactory;
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.systemui.shared.recents.model.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -70,7 +71,7 @@
     private RecentsModel(Context context) {
         mContext = context;
         mTaskList = new RecentTasksList(MAIN_EXECUTOR,
-                new KeyguardManagerCompat(context), ActivityManagerWrapper.getInstance());
+                new KeyguardManagerCompat(context), SystemUiProxy.INSTANCE.get(context));
 
         IconProvider iconProvider = new IconProvider(context);
         mIconCache = new TaskIconCache(context, RECENTS_MODEL_EXECUTOR, iconProvider);
@@ -95,7 +96,7 @@
      *                always called on the UI thread.
      * @return the request id associated with this call.
      */
-    public int getTasks(Consumer<ArrayList<Task>> callback) {
+    public int getTasks(Consumer<ArrayList<GroupTask>> callback) {
         return mTaskList.getTasks(false /* loadKeysOnly */, callback);
     }
 
@@ -120,9 +121,9 @@
      * @param callback Receives true if task is removed, false otherwise
      */
     public void isTaskRemoved(int taskId, Consumer<Boolean> callback) {
-        mTaskList.getTasks(true /* loadKeysOnly */, (tasks) -> {
-            for (Task task : tasks) {
-                if (task.key.id == taskId) {
+        mTaskList.getTasks(true /* loadKeysOnly */, (taskGroups) -> {
+            for (GroupTask group : taskGroups) {
+                if (group.containsTask(taskId)) {
                     callback.accept(false);
                     return;
                 }
@@ -148,14 +149,15 @@
         ActivityManager.RunningTaskInfo runningTask =
                 ActivityManagerWrapper.getInstance().getRunningTask();
         int runningTaskId = runningTask != null ? runningTask.id : -1;
-        mTaskList.getTaskKeys(mThumbnailCache.getCacheSize(), tasks -> {
-            for (Task task : tasks) {
-                if (task.key.id == runningTaskId) {
+        mTaskList.getTaskKeys(mThumbnailCache.getCacheSize(), taskGroups -> {
+            for (GroupTask group : taskGroups) {
+                if (group.containsTask(runningTaskId)) {
                     // Skip the running task, it's not going to have an up-to-date snapshot by the
                     // time the user next enters overview
                     continue;
                 }
-                mThumbnailCache.updateThumbnailInCache(task);
+                mThumbnailCache.updateThumbnailInCache(group.task1);
+                mThumbnailCache.updateThumbnailInCache(group.task2);
             }
         });
     }
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index b6f9d58..a875b69 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -15,6 +15,8 @@
  */
 package com.android.quickstep;
 
+import static android.app.ActivityManager.RECENT_IGNORE_UNAVAILABLE;
+
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.app.PendingIntent;
@@ -48,6 +50,9 @@
 import com.android.wm.shell.onehanded.IOneHanded;
 import com.android.wm.shell.pip.IPip;
 import com.android.wm.shell.pip.IPipAnimationListener;
+import com.android.wm.shell.recents.IRecentTasks;
+import com.android.wm.shell.recents.IRecentTasksListener;
+import com.android.wm.shell.util.GroupedRecentTaskInfo;
 import com.android.wm.shell.splitscreen.ISplitScreen;
 import com.android.wm.shell.splitscreen.ISplitScreenListener;
 import com.android.wm.shell.startingsurface.IStartingWindow;
@@ -55,6 +60,7 @@
 import com.android.wm.shell.transition.IShellTransitions;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 
 /**
  * Holds the reference to SystemUI.
@@ -73,6 +79,7 @@
     private IOneHanded mOneHanded;
     private IShellTransitions mShellTransitions;
     private IStartingWindow mStartingWindow;
+    private IRecentTasks mRecentTasks;
     private final DeathRecipient mSystemUiProxyDeathRecipient = () -> {
         MAIN_EXECUTOR.execute(() -> clearProxy());
     };
@@ -83,6 +90,7 @@
     private ISplitScreenListener mPendingSplitScreenListener;
     private IStartingWindowListener mPendingStartingWindowListener;
     private ISmartspaceCallback mPendingSmartspaceCallback;
+    private IRecentTasksListener mPendingRecentTasksListener;
     private final ArrayList<RemoteTransitionCompat> mPendingRemoteTransitions = new ArrayList<>();
 
     // Used to dedupe calls to SystemUI
@@ -118,6 +126,17 @@
     }
 
     @Override
+    public void onImeSwitcherPressed() {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.onImeSwitcherPressed();
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call onImeSwitcherPressed", e);
+            }
+        }
+    }
+
+    @Override
     public void setHomeRotationEnabled(boolean enabled) {
         if (mSystemUiProxy != null) {
             try {
@@ -136,7 +155,7 @@
 
     public void setProxy(ISystemUiProxy proxy, IPip pip, ISplitScreen splitScreen,
             IOneHanded oneHanded, IShellTransitions shellTransitions,
-            IStartingWindow startingWindow,
+            IStartingWindow startingWindow, IRecentTasks recentTasks,
             ISmartspaceTransitionController smartSpaceTransitionController) {
         unlinkToDeath();
         mSystemUiProxy = proxy;
@@ -146,6 +165,7 @@
         mShellTransitions = shellTransitions;
         mStartingWindow = startingWindow;
         mSmartspaceTransitionController = smartSpaceTransitionController;
+        mRecentTasks = recentTasks;
         linkToDeath();
         // re-attach the listeners once missing due to setProxy has not been initialized yet.
         if (mPendingPipAnimationListener != null && mPip != null) {
@@ -168,6 +188,10 @@
             registerRemoteTransition(mPendingRemoteTransitions.get(i));
         }
         mPendingRemoteTransitions.clear();
+        if (mPendingRecentTasksListener != null && mRecentTasks != null) {
+            registerRecentTasksListener(mPendingRecentTasksListener);
+            mPendingRecentTasksListener = null;
+        }
 
         if (mPendingSetNavButtonAlpha != null) {
             mPendingSetNavButtonAlpha.run();
@@ -176,7 +200,7 @@
     }
 
     public void clearProxy() {
-        setProxy(null, null, null, null, null, null, null);
+        setProxy(null, null, null, null, null, null, null, null);
     }
 
     // TODO(141886704): Find a way to remove this
@@ -748,7 +772,6 @@
         }
     }
 
-
     //
     // SmartSpace transitions
     //
@@ -764,4 +787,43 @@
             mPendingSmartspaceCallback = callback;
         }
     }
+
+    //
+    // Recents
+    //
+
+    public void registerRecentTasksListener(IRecentTasksListener listener) {
+        if (mRecentTasks != null) {
+            try {
+                mRecentTasks.registerRecentTasksListener(listener);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call registerRecentTasksListener", e);
+            }
+        } else {
+            mPendingRecentTasksListener = listener;
+        }
+    }
+
+    public void unregisterRecentTasksListener(IRecentTasksListener listener) {
+        if (mRecentTasks != null) {
+            try {
+                mRecentTasks.unregisterRecentTasksListener(listener);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call unregisterRecentTasksListener");
+            }
+        }
+        mPendingRecentTasksListener = null;
+    }
+
+    public ArrayList<GroupedRecentTaskInfo> getRecentTasks(int numTasks, int userId) {
+        if (mRecentTasks != null) {
+            try {
+                return new ArrayList<>(Arrays.asList(mRecentTasks.getRecentTasks(numTasks,
+                        RECENT_IGNORE_UNAVAILABLE, userId)));
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call getRecentTasks", e);
+            }
+        }
+        return new ArrayList<>();
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
index a8a0219..eaa43cf 100644
--- a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
+++ b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
@@ -104,6 +104,9 @@
      * Synchronously fetches the thumbnail for the given {@param task} and puts it in the cache.
      */
     public void updateThumbnailInCache(Task task) {
+        if (task == null) {
+            return;
+        }
         Preconditions.assertUIThread();
         // Fetch the thumbnail for this task and put it in the cache
         if (task.thumbnail == null) {
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index b56c1b7..c037ac4 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -26,6 +26,7 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.quickstep.GestureState.DEFAULT_STATE;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
+import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_RECENT_TASKS;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_ONE_HANDED;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_PIP;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_SHELL_TRANSITIONS;
@@ -113,6 +114,7 @@
 import com.android.systemui.shared.tracing.ProtoTraceable;
 import com.android.wm.shell.onehanded.IOneHanded;
 import com.android.wm.shell.pip.IPip;
+import com.android.wm.shell.recents.IRecentTasks;
 import com.android.wm.shell.splitscreen.ISplitScreen;
 import com.android.wm.shell.startingsurface.IStartingWindow;
 import com.android.wm.shell.transition.IShellTransitions;
@@ -171,9 +173,11 @@
             ISmartspaceTransitionController smartspaceTransitionController =
                     ISmartspaceTransitionController.Stub.asInterface(
                             bundle.getBinder(KEY_EXTRA_SMARTSPACE_TRANSITION_CONTROLLER));
+            IRecentTasks recentTasks = IRecentTasks.Stub.asInterface(
+                    bundle.getBinder(KEY_EXTRA_RECENT_TASKS));
             MAIN_EXECUTOR.execute(() -> {
                 SystemUiProxy.INSTANCE.get(TouchInteractionService.this).setProxy(proxy, pip,
-                        splitscreen, onehanded, shellTransitions, startingWindow,
+                        splitscreen, onehanded, shellTransitions, startingWindow, recentTasks,
                         smartspaceTransitionController);
                 TouchInteractionService.this.initInputMonitor();
                 preloadOverview(true /* fromInit */);
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index 7e8b83e..48315d0 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -42,6 +42,7 @@
 import com.android.quickstep.views.OverviewActionsView;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
+import com.android.systemui.shared.recents.model.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.Task.TaskKey;
 
@@ -153,29 +154,31 @@
     }
 
     @Override
-    protected void applyLoadPlan(ArrayList<Task> tasks) {
+    protected void applyLoadPlan(ArrayList<GroupTask> taskGroups) {
         // When quick-switching on 3p-launcher, we add a "stub" tile corresponding to Launcher
         // as well. This tile is never shown as we have setCurrentTaskHidden, but allows use to
         // track the index of the next task appropriately, as if we are switching on any other app.
         // TODO(b/195607777) Confirm home task info is front-most task and not mixed in with others
         int runningTaskId = getTaskIdsForRunningTaskView()[0];
-        if (mHomeTaskInfo != null && mHomeTaskInfo.taskId == runningTaskId && !tasks.isEmpty()) {
+        if (mHomeTaskInfo != null && mHomeTaskInfo.taskId == runningTaskId && !taskGroups.isEmpty()) {
             // Check if the task list has running task
             boolean found = false;
-            for (Task t : tasks) {
-                if (t.key.id == runningTaskId) {
+            for (GroupTask group : taskGroups) {
+                if (group.containsTask(runningTaskId)) {
                     found = true;
                     break;
                 }
             }
             if (!found) {
-                ArrayList<Task> newList = new ArrayList<>(tasks.size() + 1);
-                newList.addAll(tasks);
-                newList.add(Task.from(new TaskKey(mHomeTaskInfo), mHomeTaskInfo, false));
-                tasks = newList;
+                ArrayList<GroupTask> newList = new ArrayList<>(taskGroups.size() + 1);
+                newList.addAll(taskGroups);
+                newList.add(new GroupTask(
+                        Task.from(new TaskKey(mHomeTaskInfo), mHomeTaskInfo, false),
+                        null));
+                taskGroups = newList;
             }
         }
-        super.applyLoadPlan(tasks);
+        super.applyLoadPlan(taskGroups);
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/interaction/BackGestureTutorialController.java b/quickstep/src/com/android/quickstep/interaction/BackGestureTutorialController.java
index fb6cd8a..49d8203 100644
--- a/quickstep/src/com/android/quickstep/interaction/BackGestureTutorialController.java
+++ b/quickstep/src/com/android/quickstep/interaction/BackGestureTutorialController.java
@@ -43,6 +43,13 @@
     }
 
     @Override
+    public Integer getSuccessFeedbackSubtitle() {
+        return mTutorialFragment.isAtFinalStep()
+                ? R.string.back_gesture_feedback_complete_without_follow_up
+                : R.string.back_gesture_feedback_complete_with_overview_follow_up;
+    }
+
+    @Override
     protected int getMockAppTaskLayoutResId() {
         return getMockAppTaskCurrentPageLayoutResId();
     }
@@ -85,10 +92,7 @@
             case BACK_COMPLETED_FROM_RIGHT:
                 mTutorialFragment.releaseFeedbackAnimation();
                 updateFakeAppTaskViewLayout(getMockAppTaskPreviousPageLayoutResId());
-                int subtitleResId = mTutorialFragment.isAtFinalStep()
-                        ? R.string.back_gesture_feedback_complete_without_follow_up
-                        : R.string.back_gesture_feedback_complete_with_overview_follow_up;
-                showFeedback(subtitleResId, true);
+                showSuccessFeedback();
                 break;
             case BACK_CANCELLED_FROM_LEFT:
             case BACK_CANCELLED_FROM_RIGHT:
diff --git a/quickstep/src/com/android/quickstep/interaction/GestureSandboxActivity.java b/quickstep/src/com/android/quickstep/interaction/GestureSandboxActivity.java
index 7fb7d29..c2524b1 100644
--- a/quickstep/src/com/android/quickstep/interaction/GestureSandboxActivity.java
+++ b/quickstep/src/com/android/quickstep/interaction/GestureSandboxActivity.java
@@ -35,10 +35,9 @@
 /** Shows the gesture interactive sandbox in full screen mode. */
 public class GestureSandboxActivity extends FragmentActivity {
 
-    private static final String LOG_TAG = "GestureSandboxActivity";
-
     private static final String KEY_TUTORIAL_STEPS = "tutorial_steps";
     private static final String KEY_CURRENT_STEP = "current_step";
+    private static final String KEY_GESTURE_COMPLETE = "gesture_complete";
 
     private TutorialType[] mTutorialSteps;
     private TutorialType mCurrentTutorialStep;
@@ -56,7 +55,8 @@
         Bundle args = savedInstanceState == null ? getIntent().getExtras() : savedInstanceState;
         mTutorialSteps = getTutorialSteps(args);
         mCurrentTutorialStep = mTutorialSteps[mCurrentStep - 1];
-        mFragment = TutorialFragment.newInstance(mCurrentTutorialStep);
+        mFragment = TutorialFragment.newInstance(
+                mCurrentTutorialStep, args.getBoolean(KEY_GESTURE_COMPLETE, false));
         getSupportFragmentManager().beginTransaction()
                 .add(R.id.gesture_tutorial_fragment_container, mFragment)
                 .commit();
@@ -87,6 +87,7 @@
     protected void onSaveInstanceState(@NonNull Bundle savedInstanceState) {
         savedInstanceState.putStringArray(KEY_TUTORIAL_STEPS, getTutorialStepNames());
         savedInstanceState.putInt(KEY_CURRENT_STEP, mCurrentStep);
+        savedInstanceState.putBoolean(KEY_GESTURE_COMPLETE, mFragment.isGestureComplete());
         super.onSaveInstanceState(savedInstanceState);
     }
 
@@ -121,7 +122,7 @@
             return;
         }
         mCurrentTutorialStep = mTutorialSteps[mCurrentStep];
-        mFragment = TutorialFragment.newInstance(mCurrentTutorialStep);
+        mFragment = TutorialFragment.newInstance(mCurrentTutorialStep, false);
         getSupportFragmentManager().beginTransaction()
                 .replace(R.id.gesture_tutorial_fragment_container, mFragment)
                 .runOnCommit(() -> mFragment.onAttachedToWindow())
diff --git a/quickstep/src/com/android/quickstep/interaction/HomeGestureTutorialController.java b/quickstep/src/com/android/quickstep/interaction/HomeGestureTutorialController.java
index bbb22e6..0bc3691 100644
--- a/quickstep/src/com/android/quickstep/interaction/HomeGestureTutorialController.java
+++ b/quickstep/src/com/android/quickstep/interaction/HomeGestureTutorialController.java
@@ -42,6 +42,13 @@
     }
 
     @Override
+    public Integer getSuccessFeedbackSubtitle() {
+        return mTutorialFragment.isAtFinalStep()
+                ? R.string.home_gesture_feedback_complete_without_follow_up
+                : R.string.home_gesture_feedback_complete_with_follow_up;
+    }
+
+    @Override
     protected int getMockAppTaskLayoutResId() {
         return mTutorialFragment.isLargeScreen()
                 ? R.layout.gesture_tutorial_foldable_mock_webpage
@@ -84,10 +91,7 @@
                     case HOME_GESTURE_COMPLETED: {
                         mTutorialFragment.releaseFeedbackAnimation();
                         animateFakeTaskViewHome(finalVelocity, null);
-                        int subtitleResId = mTutorialFragment.isAtFinalStep()
-                                ? R.string.home_gesture_feedback_complete_without_follow_up
-                                : R.string.home_gesture_feedback_complete_with_follow_up;
-                        showFeedback(subtitleResId, true);
+                        showSuccessFeedback();
                         break;
                     }
                     case HOME_NOT_STARTED_TOO_FAR_FROM_EDGE:
diff --git a/quickstep/src/com/android/quickstep/interaction/OverviewGestureTutorialController.java b/quickstep/src/com/android/quickstep/interaction/OverviewGestureTutorialController.java
index 0fea0d7..f308f27 100644
--- a/quickstep/src/com/android/quickstep/interaction/OverviewGestureTutorialController.java
+++ b/quickstep/src/com/android/quickstep/interaction/OverviewGestureTutorialController.java
@@ -52,6 +52,13 @@
     }
 
     @Override
+    public Integer getSuccessFeedbackSubtitle() {
+        return mTutorialFragment.getNumSteps() > 1 && mTutorialFragment.isAtFinalStep()
+                ? R.string.overview_gesture_feedback_complete_with_follow_up
+                : R.string.overview_gesture_feedback_complete_without_follow_up;
+    }
+
+    @Override
     protected int getMockAppTaskLayoutResId() {
         return mTutorialFragment.isLargeScreen()
                 ? R.layout.gesture_tutorial_foldable_mock_conversation_list
@@ -106,11 +113,7 @@
                         mTutorialFragment.releaseFeedbackAnimation();
                         animateTaskViewToOverview();
                         onMotionPaused(true /*arbitrary value*/);
-                        int subtitleResId = mTutorialFragment.getNumSteps() > 1
-                                && mTutorialFragment.isAtFinalStep()
-                                ? R.string.overview_gesture_feedback_complete_with_follow_up
-                                : R.string.overview_gesture_feedback_complete_without_follow_up;
-                        showFeedback(subtitleResId, true);
+                        showSuccessFeedback();
                         break;
                     case HOME_OR_OVERVIEW_NOT_STARTED_WRONG_SWIPE_DIRECTION:
                     case HOME_OR_OVERVIEW_CANCELLED:
diff --git a/quickstep/src/com/android/quickstep/interaction/TutorialController.java b/quickstep/src/com/android/quickstep/interaction/TutorialController.java
index 81e18f6..4145393 100644
--- a/quickstep/src/com/android/quickstep/interaction/TutorialController.java
+++ b/quickstep/src/com/android/quickstep/interaction/TutorialController.java
@@ -223,6 +223,11 @@
         return null;
     }
 
+    @StringRes
+    public Integer getSuccessFeedbackSubtitle() {
+        return null;
+    }
+
     void showFeedback() {
         if (mGestureCompleted) {
             mFeedbackView.setTranslationY(0);
@@ -236,6 +241,13 @@
     }
 
     /**
+     * Show feedback reflecting a successful gesture attempt.
+     **/
+    void showSuccessFeedback() {
+        showFeedback(getSuccessFeedbackSubtitle(), true);
+    }
+
+    /**
      * Show feedback reflecting a failed gesture attempt.
      *
      * @param subtitleResId Resource of the text to display.
diff --git a/quickstep/src/com/android/quickstep/interaction/TutorialFragment.java b/quickstep/src/com/android/quickstep/interaction/TutorialFragment.java
index 89be1a6..2fd7cde 100644
--- a/quickstep/src/com/android/quickstep/interaction/TutorialFragment.java
+++ b/quickstep/src/com/android/quickstep/interaction/TutorialFragment.java
@@ -49,8 +49,10 @@
 
     private static final String LOG_TAG = "TutorialFragment";
     static final String KEY_TUTORIAL_TYPE = "tutorial_type";
+    static final String KEY_GESTURE_COMPLETE = "gesture_complete";
 
     TutorialType mTutorialType;
+    boolean mGestureComplete = false;
     @Nullable TutorialController mTutorialController = null;
     RootSandboxLayout mRootView;
     View mFingerDotView;
@@ -67,7 +69,7 @@
 
     private boolean mIsLargeScreen;
 
-    public static TutorialFragment newInstance(TutorialType tutorialType) {
+    public static TutorialFragment newInstance(TutorialType tutorialType, boolean gestureComplete) {
         TutorialFragment fragment = getFragmentForTutorialType(tutorialType);
         if (fragment == null) {
             fragment = new BackGestureTutorialFragment();
@@ -76,6 +78,7 @@
 
         Bundle args = new Bundle();
         args.putSerializable(KEY_TUTORIAL_TYPE, tutorialType);
+        args.putBoolean(KEY_GESTURE_COMPLETE, gestureComplete);
         fragment.setArguments(args);
         return fragment;
     }
@@ -132,6 +135,7 @@
         super.onCreate(savedInstanceState);
         Bundle args = savedInstanceState != null ? savedInstanceState : getArguments();
         mTutorialType = (TutorialType) args.getSerializable(KEY_TUTORIAL_TYPE);
+        mGestureComplete = args.getBoolean(KEY_GESTURE_COMPLETE, false);
         mEdgeBackGestureHandler = new EdgeBackGestureHandler(getContext());
         mNavBarGestureHandler = new NavBarGestureHandler(getContext());
 
@@ -186,11 +190,13 @@
     }
 
     void initializeFeedbackVideoView() {
-        if (!updateFeedbackAnimation()) {
+        if (!updateFeedbackAnimation() || mTutorialController == null) {
             return;
         }
 
-        if (!mIntroductionShown && mTutorialController != null) {
+        if (isGestureComplete()) {
+            mTutorialController.showSuccessFeedback();
+        } else if (!mIntroductionShown) {
             Integer introTileStringResId = mTutorialController.getIntroductionTitle();
             Integer introSubtitleResId = mTutorialController.getIntroductionSubtitle();
             if (introTileStringResId != null && introSubtitleResId != null) {
@@ -372,6 +378,11 @@
         return getCurrentStep() == getNumSteps();
     }
 
+    boolean isGestureComplete() {
+        return mGestureComplete
+                || (mTutorialController != null && mTutorialController.isGestureCompleted());
+    }
+
     @Nullable
     private GestureSandboxActivity getGestureSandboxActivity() {
         Context context = getContext();
diff --git a/quickstep/src/com/android/quickstep/views/OverviewActionsView.java b/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
index 76d3591..5a8cbc1 100644
--- a/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
+++ b/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
@@ -20,6 +20,7 @@
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.util.AttributeSet;
+import android.util.Log;
 import android.view.View;
 import android.view.View.OnClickListener;
 import android.widget.FrameLayout;
@@ -30,6 +31,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.R;
+import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.launcher3.util.MultiValueAlpha.AlphaProperty;
 import com.android.quickstep.SysUINavigationMode;
@@ -110,6 +112,11 @@
     protected void onFinishInflate() {
         super.onFinishInflate();
         findViewById(R.id.action_screenshot).setOnClickListener(this);
+        if (TestProtocol.sDebugTracing) {
+            Log.d(TestProtocol.NO_SCREENSHOT, "Inflated OverviewActionsView and added screenshot"
+                    + " listener.");
+        }
+
         mSplitButton = findViewById(R.id.action_split);
         mSplitButton.setOnClickListener(this);
     }
@@ -125,6 +132,11 @@
 
     @Override
     public void onClick(View view) {
+        if (TestProtocol.sDebugTracing) {
+            Log.d(TestProtocol.NO_SCREENSHOT, "OverviewActionsView - onClick"
+                    + " callbacks: " + mCallbacks + "  view id: " + view.getId() + " "
+                    + " is screenshot? " + (view.getId() == R.id.action_screenshot));
+        }
         if (mCallbacks == null) {
             return;
         }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 96c5561..b5238c6 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -167,6 +167,7 @@
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.util.VibratorWrapper;
 import com.android.systemui.plugins.ResourceProvider;
+import com.android.systemui.shared.recents.model.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -1296,13 +1297,13 @@
         updateGridProperties();
     }
 
-    protected void applyLoadPlan(ArrayList<Task> tasks) {
+    protected void applyLoadPlan(ArrayList<GroupTask> taskGroups) {
         if (mPendingAnimation != null) {
-            mPendingAnimation.addEndListener(success -> applyLoadPlan(tasks));
+            mPendingAnimation.addEndListener(success -> applyLoadPlan(taskGroups));
             return;
         }
 
-        if (tasks == null || tasks.isEmpty()) {
+        if (taskGroups == null || taskGroups.isEmpty()) {
             removeTasksViewsAndClearAllButton();
             onTaskStackUpdated();
             return;
@@ -1324,10 +1325,11 @@
                 LauncherSplitScreenListener.INSTANCE.getNoCreate().getPersistentSplitIds();
         int requiredGroupTaskViews = splitTaskIds.length / 2;
 
+        // TODO(b/202740477): Update once grouped tasks are returned
         // Subtract half the number of split tasks and not total number because we've already
         // added a GroupedTaskView when swipe up gesture happens.
         // This will need to change if we start showing GroupedTaskViews during swipe up from home
-        int requiredTaskViewCount = tasks.size() - requiredGroupTaskViews;
+        int requiredTaskViewCount = taskGroups.size() - requiredGroupTaskViews;
 
         if (getTaskViewCount() != requiredTaskViewCount) {
             if (indexOfChild(mClearAllButton) != -1) {
@@ -1360,11 +1362,12 @@
                 + " runningTaskViewId: " + mRunningTaskViewId
                 + " forTaskView: " + getTaskViewFromTaskViewId(mRunningTaskViewId));
 
-        for (int taskViewIndex = requiredTaskViewCount - 1, taskDataIndex = tasks.size() - 1;
+        for (int taskViewIndex = requiredTaskViewCount - 1, taskDataIndex = taskGroups.size() - 1;
                 taskViewIndex >= 0;
                 taskViewIndex--, taskDataIndex--) {
             final int pageIndex = requiredTaskViewCount - taskViewIndex - 1;
-            final Task task = tasks.get(taskDataIndex);
+            // TODO(b/202740477): Temporary assumption, to be updated once groups are actually used
+            final Task task = taskGroups.get(taskDataIndex).task1;
             final TaskView taskView = (TaskView) getChildAt(pageIndex);
             if (taskView instanceof GroupedTaskView) {
                 Task leftTop;
@@ -1372,11 +1375,11 @@
                 if (task.key.id == splitTaskIds[0]) {
                     leftTop = task;
                     taskDataIndex--;
-                    rightBottom = tasks.get(taskDataIndex);
+                    rightBottom = taskGroups.get(taskDataIndex).task1;
                 } else {
                     rightBottom = task;
                     taskDataIndex--;
-                    leftTop = tasks.get(taskDataIndex);
+                    leftTop = taskGroups.get(taskDataIndex).task1;
                 }
                 ((GroupedTaskView) taskView).bind(leftTop, rightBottom, mOrientationState,
                         mSplitBoundsConfig);
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index 77ac373..2690a2a 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -270,15 +270,9 @@
         BaseDragLayer.LayoutParams params = (BaseDragLayer.LayoutParams) getLayoutParams();
         int padding = getResources()
                 .getDimensionPixelSize(R.dimen.task_menu_vertical_padding);
-        if (deviceProfile.overviewShowAsGrid) {
-            // TODO(b/193432925) temporary so it doesn't look terrible on large screen
-            params.width =
-                    getContext().getResources().getDimensionPixelSize(R.dimen.task_menu_width_grid);
-        } else {
-            params.width = orientationHandler
-                    .getTaskMenuWidth(taskContainer.getThumbnailView(),
-                            deviceProfile) - (2 * padding);
-        }
+        params.width = orientationHandler
+                .getTaskMenuWidth(taskContainer.getThumbnailView(),
+                        deviceProfile) - (2 * padding);
         // Gravity set to Left instead of Start as sTempRect.left measures Left distance not Start
         params.gravity = Gravity.LEFT;
         setLayoutParams(params);
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
index 9b86c73..39a6fc4 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
@@ -1,14 +1,150 @@
 package com.android.quickstep.views
 
-import android.util.Log
+import android.animation.AnimatorSet
+import android.animation.ObjectAnimator
+import android.content.Context
+import android.graphics.Rect
+import android.graphics.drawable.ShapeDrawable
+import android.graphics.drawable.shapes.RectShape
+import android.util.AttributeSet
+import android.view.MotionEvent
+import android.view.View
+import android.view.ViewGroup
+import android.widget.LinearLayout
+import com.android.launcher3.BaseDraggingActivity
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.R
+import com.android.launcher3.popup.ArrowPopup
+import com.android.launcher3.popup.SystemShortcut
+import com.android.launcher3.util.Themes
+import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
 
-// TODO(http://b/193432925)
-class TaskMenuViewWithArrow {
+class TaskMenuViewWithArrow<T : BaseDraggingActivity> : ArrowPopup<T> {
     companion object {
         const val TAG = "TaskMenuViewWithArrow"
 
-        fun logSomething() {
-            Log.d(TAG, "It worked!")
+        fun showForTask(taskContainer: TaskIdAttributeContainer): Boolean {
+            val activity = BaseDraggingActivity
+                    .fromContext<BaseDraggingActivity>(taskContainer.taskView.context)
+            val taskMenuViewWithArrow = activity.layoutInflater
+                    .inflate(R.layout.task_menu_with_arrow, activity.dragLayer, false) as TaskMenuViewWithArrow<*>
+
+            return taskMenuViewWithArrow.populateAndShowForTask(taskContainer)
         }
     }
+
+    constructor(context: Context) : super(context)
+    constructor(context: Context, attrs: AttributeSet) : super(context, attrs)
+    constructor(context: Context, attrs: AttributeSet, defStyleAttr: Int) : super(context, attrs, defStyleAttr)
+
+    init {
+        clipToOutline = true
+    }
+
+    private val menuWidth = context.resources.getDimensionPixelSize(R.dimen.task_menu_width_grid)
+
+    private lateinit var taskView: TaskView
+    private lateinit var optionLayout: LinearLayout
+    private lateinit var taskContainer: TaskIdAttributeContainer
+
+    override fun isOfType(type: Int): Boolean = type and TYPE_TASK_MENU != 0
+
+    override fun getTargetObjectLocation(outPos: Rect?) {
+        popupContainer.getDescendantRectRelativeToSelf(taskView.iconView, outPos)
+    }
+
+    override fun onControllerInterceptTouchEvent(ev: MotionEvent?): Boolean {
+        if (ev?.action == MotionEvent.ACTION_DOWN) {
+            if (!popupContainer.isEventOverView(this, ev)) {
+                close(true)
+                return true
+            }
+        }
+        return false
+    }
+
+    override fun onFinishInflate() {
+        super.onFinishInflate()
+        optionLayout = findViewById(R.id.menu_option_layout)
+    }
+
+    private fun populateAndShowForTask(taskContainer: TaskIdAttributeContainer): Boolean {
+        if (isAttachedToWindow) {
+            return false
+        }
+
+        taskView = taskContainer.taskView
+        this.taskContainer = taskContainer
+        if (!populateMenu()) return false
+        show()
+        return true
+    }
+
+    /** @return true if successfully able to populate task view menu, false otherwise
+     */
+    private fun populateMenu(): Boolean {
+        // Icon may not be loaded
+        if (taskContainer.task.icon == null) return false
+
+        addMenuOptions()
+        return true
+    }
+
+    private fun addMenuOptions() {
+        // Add the options
+        TaskOverlayFactory
+            .getEnabledShortcuts(taskView, mActivityContext.deviceProfile, taskContainer)
+            .forEach { this.addMenuOption(it) }
+
+        // Add the spaces between items
+        val divider = ShapeDrawable(RectShape())
+        divider.paint.color = resources.getColor(android.R.color.transparent)
+        val dividerSpacing = resources.getDimension(R.dimen.task_menu_spacing).toInt()
+        optionLayout.showDividers = SHOW_DIVIDER_MIDDLE
+
+        // Set the orientation, which makes the menu show
+        val recentsView: RecentsView<*, *> = mActivityContext.getOverviewPanel()
+        val orientationHandler = recentsView.pagedOrientationHandler
+        val deviceProfile: DeviceProfile = mActivityContext.deviceProfile
+        orientationHandler.setTaskOptionsMenuLayoutOrientation(
+            deviceProfile,
+            optionLayout,
+            dividerSpacing,
+            divider
+        )
+    }
+
+    private fun addMenuOption(menuOption: SystemShortcut<*>) {
+        val menuOptionView = mActivityContext.layoutInflater.inflate(
+            R.layout.task_view_menu_option, this, false
+        ) as LinearLayout
+        menuOption.setIconAndLabelFor(
+            menuOptionView.findViewById(R.id.icon),
+            menuOptionView.findViewById(R.id.text)
+        )
+        val lp = menuOptionView.layoutParams as LayoutParams
+        lp.width = menuWidth
+        menuOptionView.setOnClickListener { view: View? -> menuOption.onClick(view) }
+        optionLayout.addView(menuOptionView)
+    }
+
+    override fun assignMarginsAndBackgrounds(viewGroup: ViewGroup) {
+        assignMarginsAndBackgrounds(this, Themes.getAttrColor(context, com.android.internal.R.attr.colorSurface))
+    }
+
+    override fun onCreateOpenAnimation(anim: AnimatorSet) {
+        anim.play(
+            ObjectAnimator.ofFloat(
+                taskContainer.thumbnailView, TaskThumbnailView.DIM_ALPHA,
+                TaskView.MAX_PAGE_SCRIM_ALPHA
+            )
+        )
+    }
+
+    override fun onCreateCloseAnimation(anim: AnimatorSet) {
+        anim.play(
+            ObjectAnimator.ofFloat(taskContainer.thumbnailView, TaskThumbnailView.DIM_ALPHA, 0f)
+        )
+    }
 }
\ No newline at end of file
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index f88b243..df1817e 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -824,9 +824,11 @@
     }
 
     protected boolean showTaskMenuWithContainer(IconView iconView) {
-        // TODO(http://b/193432925)
-        if (DEBUG) TaskMenuViewWithArrow.Companion.logSomething();
-        return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        if (mActivity.getDeviceProfile().overviewShowAsGrid) {
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
+        } else {
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        }
     }
 
     protected void setIcon(IconView iconView, Drawable icon) {
diff --git a/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java b/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
index 79ddf7a..6b2d5ed 100644
--- a/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
+++ b/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
@@ -30,70 +30,78 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.util.LooperExecutor;
+import com.android.systemui.shared.recents.model.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.KeyguardManagerCompat;
+import com.android.wm.shell.util.GroupedRecentTaskInfo;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
 @SmallTest
 public class RecentTasksListTest {
 
-    private ActivityManagerWrapper mockActivityManagerWrapper;
+    @Mock
+    private SystemUiProxy mockSystemUiProxy;
 
     // Class under test
     private RecentTasksList mRecentTasksList;
 
     @Before
     public void setup() {
+        MockitoAnnotations.initMocks(this);
         LooperExecutor mockMainThreadExecutor = mock(LooperExecutor.class);
         KeyguardManagerCompat mockKeyguardManagerCompat = mock(KeyguardManagerCompat.class);
-        mockActivityManagerWrapper = mock(ActivityManagerWrapper.class);
         mRecentTasksList = new RecentTasksList(mockMainThreadExecutor, mockKeyguardManagerCompat,
-                mockActivityManagerWrapper);
+                mockSystemUiProxy);
     }
 
     @Test
-    public void onTaskRemoved_doesNotFetchTasks() {
-        mRecentTasksList.onTaskRemoved(0);
-        verify(mockActivityManagerWrapper, times(0))
-                .getRecentTasks(anyInt(), anyInt());
-    }
-
-    @Test
-    public void onTaskStackChanged_doesNotFetchTasks() {
-        mRecentTasksList.onTaskStackChanged();
-        verify(mockActivityManagerWrapper, times(0))
+    public void onRecentTasksChanged_doesNotFetchTasks() {
+        mRecentTasksList.onRecentTasksChanged();
+        verify(mockSystemUiProxy, times(0))
                 .getRecentTasks(anyInt(), anyInt());
     }
 
     @Test
     public void loadTasksInBackground_onlyKeys_noValidTaskDescription() {
-        ActivityManager.RecentTaskInfo recentTaskInfo = new ActivityManager.RecentTaskInfo();
-        when(mockActivityManagerWrapper.getRecentTasks(anyInt(), anyInt()))
-                .thenReturn(Collections.singletonList(recentTaskInfo));
+        GroupedRecentTaskInfo recentTaskInfos = new GroupedRecentTaskInfo(
+                new ActivityManager.RecentTaskInfo(), new ActivityManager.RecentTaskInfo());
+        when(mockSystemUiProxy.getRecentTasks(anyInt(), anyInt()))
+                .thenReturn(new ArrayList<>(Collections.singletonList(recentTaskInfos)));
 
-        List<Task> taskList = mRecentTasksList.loadTasksInBackground(Integer.MAX_VALUE, -1, true);
+        List<GroupTask> taskList = mRecentTasksList.loadTasksInBackground(Integer.MAX_VALUE, -1,
+                true);
 
         assertEquals(1, taskList.size());
-        assertNull(taskList.get(0).taskDescription.getLabel());
+        assertNull(taskList.get(0).task1.taskDescription.getLabel());
+        assertNull(taskList.get(0).task2.taskDescription.getLabel());
     }
 
     @Test
     public void loadTasksInBackground_moreThanKeys_hasValidTaskDescription() {
         String taskDescription = "Wheeee!";
-        ActivityManager.RecentTaskInfo recentTaskInfo = new ActivityManager.RecentTaskInfo();
-        recentTaskInfo.taskDescription = new ActivityManager.TaskDescription(taskDescription);
-        when(mockActivityManagerWrapper.getRecentTasks(anyInt(), anyInt()))
-                .thenReturn(Collections.singletonList(recentTaskInfo));
+        ActivityManager.RecentTaskInfo task1 = new ActivityManager.RecentTaskInfo();
+        task1.taskDescription = new ActivityManager.TaskDescription(taskDescription);
+        ActivityManager.RecentTaskInfo task2 = new ActivityManager.RecentTaskInfo();
+        task2.taskDescription = new ActivityManager.TaskDescription();
+        GroupedRecentTaskInfo recentTaskInfos = new GroupedRecentTaskInfo(
+                task1, task2);
+        when(mockSystemUiProxy.getRecentTasks(anyInt(), anyInt()))
+                .thenReturn(new ArrayList<>(Collections.singletonList(recentTaskInfos)));
 
-        List<Task> taskList = mRecentTasksList.loadTasksInBackground(Integer.MAX_VALUE, -1, false);
+        List<GroupTask> taskList = mRecentTasksList.loadTasksInBackground(Integer.MAX_VALUE, -1,
+                false);
 
         assertEquals(1, taskList.size());
-        assertEquals(taskDescription, taskList.get(0).taskDescription.getLabel());
+        assertEquals(taskDescription, taskList.get(0).task1.taskDescription.getLabel());
+        assertNull(taskList.get(0).task2.taskDescription.getLabel());
     }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 437a19b..ea65757 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -44,6 +44,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -292,6 +293,7 @@
 
     // TODO(b/204830798): test with all navigation modes(add @NavigationModeSwitch annotation)
     //  after the bug resolved.
+    @Ignore("b/205027405")
     @Test
     @PortraitLandscape
     @ScreenRecord
diff --git a/res/layout/all_apps.xml b/res/layout/all_apps.xml
index a34baef..7f9f63e 100644
--- a/res/layout/all_apps.xml
+++ b/res/layout/all_apps.xml
@@ -43,15 +43,6 @@
 
         <include layout="@layout/all_apps_personal_work_tabs" />
 
-        <Button
-            android:id="@+id/all_apps_button"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:layout_gravity="center"
-            android:text="@string/all_apps_label"
-            android:background="@drawable/padded_rounded_action_button"
-            android:visibility="gone"/>
-
     </com.android.launcher3.allapps.FloatingHeaderView>
 
     <include layout="@layout/search_container_all_apps" />
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index fc717c9..ce06c6e 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -329,6 +329,7 @@
         }
 
         updateWorkspaceScreensPadding();
+        updateWorkspaceWidgetsSizes();
     }
 
     private void updateWorkspaceScreensPadding() {
@@ -360,6 +361,25 @@
         }
     }
 
+    private void updateWorkspaceWidgetsSizes() {
+        int numberOfScreens = mScreenOrder.size();
+        for (int i = 0; i < numberOfScreens; i++) {
+            ShortcutAndWidgetContainer shortcutAndWidgetContainer =
+                    mWorkspaceScreens.get(mScreenOrder.get(i)).getShortcutsAndWidgets();
+            int shortcutsAndWidgetCount = shortcutAndWidgetContainer.getChildCount();
+            for (int j = 0; j < shortcutsAndWidgetCount; j++) {
+                View view = shortcutAndWidgetContainer.getChildAt(j);
+                if (view instanceof LauncherAppWidgetHostView
+                        && view.getTag() instanceof LauncherAppWidgetInfo) {
+                    LauncherAppWidgetInfo launcherAppWidgetInfo =
+                            (LauncherAppWidgetInfo) view.getTag();
+                    WidgetSizes.updateWidgetSizeRanges((LauncherAppWidgetHostView) view,
+                            mLauncher, launcherAppWidgetInfo.spanX, launcherAppWidgetInfo.spanY);
+                }
+            }
+        }
+    }
+
     /**
      * Estimates the size of an item using spans: hSpan, vSpan.
      *
diff --git a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
index 874fe80..f089551 100644
--- a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
+++ b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
@@ -42,6 +42,7 @@
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.R;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.util.PackageManagerHelper;
 
 import java.util.Arrays;
@@ -95,16 +96,15 @@
         // The type of this item
         public int viewType;
 
-        /** App-only properties */
-        // The section name of this app.  Note that there can be multiple items with different
+        // The section name of this item.  Note that there can be multiple items with different
         // sectionNames in the same section
         public String sectionName = null;
         // The row that this item shows up on
         public int rowIndex;
         // The index of this app in the row
         public int rowAppIndex;
-        // The associated AppInfo for the app
-        public AppInfo appInfo = null;
+        // The associated ItemInfoWithIcon for the item
+        public ItemInfoWithIcon itemInfo = null;
         // The index of this app not including sections
         public int appIndex = -1;
         // Search section associated to result
@@ -119,7 +119,7 @@
             item.viewType = VIEW_TYPE_ICON;
             item.position = pos;
             item.sectionName = sectionName;
-            item.appInfo = appInfo;
+            item.itemInfo = appInfo;
             item.appIndex = appIndex;
             return item;
         }
@@ -373,7 +373,7 @@
                 if (adapterProvider != null) {
                     return adapterProvider.onCreateViewHolder(mLayoutInflater, parent, viewType);
                 }
-                throw new RuntimeException("Unexpected view type");
+                throw new RuntimeException("Unexpected view type" + viewType);
         }
     }
 
@@ -382,10 +382,13 @@
         switch (holder.getItemViewType()) {
             case VIEW_TYPE_ICON:
                 AdapterItem adapterItem = mApps.getAdapterItems().get(position);
-                AppInfo info = adapterItem.appInfo;
                 BubbleTextView icon = (BubbleTextView) holder.itemView;
                 icon.reset();
-                icon.applyFromApplicationInfo(info);
+                if (adapterItem.itemInfo instanceof AppInfo) {
+                    icon.applyFromApplicationInfo((AppInfo) adapterItem.itemInfo);
+                } else {
+                    icon.applyFromItemInfoWithIcon(adapterItem.itemInfo);
+                }
                 break;
             case VIEW_TYPE_EMPTY_SEARCH:
                 TextView emptyViewText = (TextView) holder.itemView;
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index 3ca0303..85ee636 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -47,7 +47,6 @@
         ValueAnimator.AnimatorUpdateListener, PluginListener<AllAppsRow>, Insettable,
         OnHeightUpdatedListener {
 
-    private static final long ALL_APPS_CONTENT_ANIM_DURATION = 150;
     private final Rect mRVClip = new Rect(0, 0, Integer.MAX_VALUE, Integer.MAX_VALUE);
     private final Rect mHeaderClip = new Rect(0, 0, Integer.MAX_VALUE, Integer.MAX_VALUE);
     private final ValueAnimator mAnimator = ValueAnimator.ofInt(0, 0);
@@ -110,13 +109,6 @@
     private FloatingHeaderRow[] mAllRows = FloatingHeaderRow.NO_ROWS;
 
 
-    // members for handling suggestion state
-    private final ValueAnimator mAllAppsContentAnimator = ValueAnimator.ofFloat(0, 0);
-    private View mAllAppsButton;
-    private int mAllAppsContentFadeInOffset;
-    private boolean mInSuggestionMode = false;
-
-
     public FloatingHeaderView(@NonNull Context context) {
         this(context, null);
     }
@@ -127,20 +119,12 @@
                 .getDimensionPixelSize(R.dimen.all_apps_header_top_padding);
         mHeaderProtectionSupported = context.getResources().getBoolean(
                 R.bool.config_header_protection_supported);
-        mAllAppsContentFadeInOffset = context.getResources()
-                .getDimensionPixelSize(R.dimen.all_apps_content_fade_in_offset);
-        mAllAppsContentAnimator.setDuration(ALL_APPS_CONTENT_ANIM_DURATION);
-        mAllAppsContentAnimator.addUpdateListener(this::onAllAppsContentAnimationUpdate);
     }
 
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
         mTabLayout = findViewById(R.id.tabs);
-        mAllAppsButton = findViewById(R.id.all_apps_button);
-        if (mAllAppsButton != null) {
-            mAllAppsButton.setOnClickListener(this::onAllAppsButtonClicked);
-        }
 
         // Find all floating header rows.
         ArrayList<FloatingHeaderRow> rows = new ArrayList<>();
@@ -329,7 +313,6 @@
         }
 
         mTabLayout.setTranslationY(mTranslationY);
-        setSuggestionMode(false);
 
         int clipHeight = mHeaderTopPadding - getPaddingBottom();
         mRVClip.top = mTabsHidden ? clipHeight : 0;
@@ -365,7 +348,6 @@
             mTranslationY = 0;
             applyVerticalMove();
         }
-        setSuggestionMode(false);
         mHeaderCollapsed = false;
         mSnappedScrolledY = -mMaxTranslation;
         mCurrentRV.scrollToTop();
@@ -461,38 +443,6 @@
         }
         return Math.max(getHeight() - getPaddingTop() + mTranslationY, 0);
     }
-
-    /**
-     * When suggestion mode is enabled, hides AllApps content view and shows AllApps button.
-     */
-    public void setSuggestionMode(boolean isSuggestMode) {
-        if (mInSuggestionMode == isSuggestMode || mAllAppsButton == null) return;
-        if (!FeatureFlags.ENABLE_ONE_SEARCH.get()) return;
-        AllAppsContainerView allApps = (AllAppsContainerView) getParent();
-        mInSuggestionMode = isSuggestMode;
-        if (isSuggestMode) {
-            mTabLayout.setVisibility(GONE);
-            mAllAppsButton.setVisibility(VISIBLE);
-            allApps.getContentView().setVisibility(GONE);
-        } else {
-            mTabLayout.setVisibility(mTabsHidden ? GONE : VISIBLE);
-            mAllAppsButton.setVisibility(GONE);
-            allApps.getContentView().setVisibility(VISIBLE);
-        }
-    }
-
-    private void onAllAppsButtonClicked(View view) {
-        setSuggestionMode(false);
-        mAllAppsContentAnimator.start();
-    }
-
-    private void onAllAppsContentAnimationUpdate(ValueAnimator valueAnimator) {
-        float prog = valueAnimator.getAnimatedFraction();
-        View allAppsList = ((AllAppsContainerView) getParent()).getContentView();
-        allAppsList.setAlpha(255 * prog);
-        allAppsList.setTranslationY((1 - prog) * mAllAppsContentFadeInOffset);
-    }
-
 }
 
 
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index 2230914..5a1e4bf 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -234,7 +234,7 @@
      * @param backgroundColor When Color.TRANSPARENT, we get color from {@link #mColorIds}.
      *                        Otherwise, we will use this color for all child views.
      */
-    private void assignMarginsAndBackgrounds(ViewGroup viewGroup, int backgroundColor) {
+    protected void assignMarginsAndBackgrounds(ViewGroup viewGroup, int backgroundColor) {
         int[] colors = null;
         if (backgroundColor == Color.TRANSPARENT) {
             // Lazily get the colors so they match the current wallpaper colors.
@@ -445,7 +445,7 @@
         animateOpen();
     }
 
-    private void setupForDisplay() {
+    protected void setupForDisplay() {
         setVisibility(View.INVISIBLE);
         mIsOpen = true;
         getPopupContainer().addView(this);
@@ -482,7 +482,7 @@
         mArrow.setVisibility(show && shouldAddArrow() ? VISIBLE : INVISIBLE);
     }
 
-    private void addArrow() {
+    protected void addArrow() {
         getPopupContainer().addView(mArrow);
         mArrow.setX(getX() + getArrowLeft());
 
@@ -686,12 +686,13 @@
         return getChildCount() > 0 ? getChildAt(0) : this;
     }
 
-    private void animateOpen() {
+    protected void animateOpen() {
         setVisibility(View.VISIBLE);
 
         mOpenCloseAnimator = getOpenCloseAnimator(true, OPEN_DURATION, OPEN_FADE_START_DELAY,
                 OPEN_FADE_DURATION, OPEN_CHILD_FADE_START_DELAY, OPEN_CHILD_FADE_DURATION,
                 DECELERATED_EASE);
+        onCreateOpenAnimation(mOpenCloseAnimator);
         mOpenCloseAnimator.addListener(new AnimatorListenerAdapter() {
             @Override
             public void onAnimationEnd(Animator animation) {
@@ -785,6 +786,11 @@
     }
 
     /**
+     * Called when creating the open transition allowing subclass can add additional animations.
+     */
+    protected void onCreateOpenAnimation(AnimatorSet anim) { }
+
+    /**
      * Called when creating the close transition allowing subclass can add additional animations.
      */
     protected void onCreateCloseAnimation(AnimatorSet anim) { }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 82163cb..b3457cd 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -346,13 +346,19 @@
     }
 
     // Cannot be used in TaplTests between a Tapl call injecting a gesture and a tapl call
-    // expecting
-    // the results of that gesture because the wait can hide flakeness.
+    // expecting the results of that gesture because the wait can hide flakeness.
     protected void waitForState(String message, Supplier<LauncherState> state) {
         waitForLauncherCondition(message,
                 launcher -> launcher.getStateManager().getCurrentStableState() == state.get());
     }
 
+    // Cannot be used in TaplTests between a Tapl call injecting a gesture and a tapl call
+    // expecting the results of that gesture because the wait can hide flakeness.
+    protected void waitForStableState(String message, Supplier<LauncherState> state) {
+        waitForLauncherCondition(message,
+                launcher -> launcher.getStateManager().isInStableState(state.get()));
+    }
+
     protected void waitForResumed(String message) {
         waitForLauncherCondition(message, launcher -> launcher.hasBeenResumed());
     }
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 4007c26..f29ac23 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -387,6 +387,7 @@
         return appIcon;
     }
 
+    @Ignore("b/205014516")
     @Test
     @PortraitLandscape
     public void testDragToFolder() throws Exception {
@@ -413,6 +414,7 @@
         folder.close();
     }
 
+    @Ignore("b/205027405")
     @Test
     @PortraitLandscape
     public void testPressBack() throws Exception {
diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 27a2375..2087bfe 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -91,9 +91,9 @@
     public void workTabExists() {
         mDevice.pressHome();
         waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForState("Launcher internal state didn't switch to Normal", () -> NORMAL);
+        waitForStableState("Launcher internal state didn't switch to Normal", () -> NORMAL);
         executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
-        waitForState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
+        waitForStableState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
         waitForLauncherCondition("Personal tab is missing",
                 launcher -> launcher.getAppsView().isPersonalTabVisible(),
                 LauncherInstrumentation.WAIT_TIME_MS);