diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/FallbackSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/FallbackSwipeHandler.java
index 888ea9c..700feef 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/FallbackSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/FallbackSwipeHandler.java
@@ -453,8 +453,10 @@
 
     @Override
     public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
-        super.onRecentsAnimationCanceled(thumbnailData);
         mStateCallback.setStateOnUiThread(STATE_HANDLER_INVALIDATED);
+
+        // Defer clearing the controller and the targets until after we've updated the state
+        super.onRecentsAnimationCanceled(thumbnailData);
     }
 
     /**
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
index 8f75c79..1b60404 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
@@ -610,10 +610,12 @@
 
     @Override
     public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
-        super.onRecentsAnimationCanceled(thumbnailData);
+        ActiveGestureLog.INSTANCE.addLog("cancelRecentsAnimation");
         mActivityInitListener.unregister();
         mStateCallback.setStateOnUiThread(STATE_GESTURE_CANCELLED | STATE_HANDLER_INVALIDATED);
-        ActiveGestureLog.INSTANCE.addLog("cancelRecentsAnimation");
+
+        // Defer clearing the controller and the targets until after we've updated the state
+        super.onRecentsAnimationCanceled(thumbnailData);
     }
 
     @Override
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
index 0f0fda9..bcaa126 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
@@ -1310,7 +1310,7 @@
     }
 
     private void dismissCurrentTask() {
-        TaskView taskView = getTaskView(getNextPage());
+        TaskView taskView = getNextPageTaskView();
         if (taskView != null) {
             dismissTask(taskView, true /*animateTaskView*/, true /*removeTask*/);
         }
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 5606ac2..b786c8b 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -35,6 +35,7 @@
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.TestHelpers;
+import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.FailureWatcher;
 import com.android.systemui.shared.system.QuickStepContract;
 
@@ -57,6 +58,8 @@
 
     static final String TAG = "QuickStepOnOffRule";
 
+    public static final int WAIT_TIME_MS = 10000;
+
     public enum Mode {
         THREE_BUTTON, TWO_BUTTON, ZERO_BUTTON, ALL
     }
@@ -118,8 +121,8 @@
                         if (mode == THREE_BUTTON || mode == ALL) {
                             evaluateWithThreeButtons();
                         }
-                    } catch (Exception e) {
-                        Log.e(TAG, "Exception", e);
+                    } catch (Throwable e) {
+                        Log.e(TAG, "Error", e);
                         throw e;
                     } finally {
                         assertTrue("Couldn't set overlay",
@@ -195,19 +198,14 @@
                                 currentSysUiNavigationMode() == expectedMode);
                     }
 
-                    for (int i = 0; i != 100; ++i) {
-                        if (mLauncher.getNavigationModel() == expectedMode) break;
-                        Thread.sleep(100);
-                    }
-                    assertTrue("Couldn't switch to " + overlayPackage,
-                            mLauncher.getNavigationModel() == expectedMode);
+                    Wait.atMost("Couldn't switch to " + overlayPackage,
+                            () -> mLauncher.getNavigationModel() == expectedMode, WAIT_TIME_MS,
+                            mLauncher);
 
-                    for (int i = 0; i != 100; ++i) {
-                        if (mLauncher.getNavigationModeMismatchError() == null) break;
-                        Thread.sleep(100);
-                    }
-                    final String error = mLauncher.getNavigationModeMismatchError();
-                    assertTrue("Switching nav mode: " + error, error == null);
+                    Wait.atMost(() -> "Switching nav mode: "
+                                    + mLauncher.getNavigationModeMismatchError(),
+                            () -> mLauncher.getNavigationModeMismatchError() == null, WAIT_TIME_MS,
+                            mLauncher);
 
                     return true;
                 }
diff --git a/src/com/android/launcher3/compat/AccessibilityManagerCompat.java b/src/com/android/launcher3/compat/AccessibilityManagerCompat.java
index d47a40e..28579c1 100644
--- a/src/com/android/launcher3/compat/AccessibilityManagerCompat.java
+++ b/src/com/android/launcher3/compat/AccessibilityManagerCompat.java
@@ -18,6 +18,7 @@
 
 import android.content.Context;
 import android.os.Bundle;
+import android.util.Log;
 import android.view.View;
 import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.AccessibilityManager;
@@ -57,6 +58,7 @@
         parcel.putInt(TestProtocol.STATE_FIELD, stateOrdinal);
 
         sendEventToTest(accessibilityManager, TestProtocol.SWITCHED_TO_STATE_MESSAGE, parcel);
+        Log.d(TestProtocol.PERMANENT_DIAG_TAG, "sendStateEventToTest: " + stateOrdinal);
     }
 
     public static void sendScrollFinishedEventToTest(Context context) {
diff --git a/src/com/android/launcher3/logging/DumpTargetWrapper.java b/src/com/android/launcher3/logging/DumpTargetWrapper.java
index 365e8f2..067bdfd 100644
--- a/src/com/android/launcher3/logging/DumpTargetWrapper.java
+++ b/src/com/android/launcher3/logging/DumpTargetWrapper.java
@@ -15,17 +15,22 @@
  */
 package com.android.launcher3.logging;
 
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
+
+import android.content.ComponentName;
 import android.os.Process;
 import android.text.TextUtils;
 
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.WorkspaceItemInfo;
 import com.android.launcher3.model.nano.LauncherDumpProto;
 import com.android.launcher3.model.nano.LauncherDumpProto.ContainerType;
 import com.android.launcher3.model.nano.LauncherDumpProto.DumpTarget;
 import com.android.launcher3.model.nano.LauncherDumpProto.ItemType;
 import com.android.launcher3.model.nano.LauncherDumpProto.UserType;
+import com.android.launcher3.util.ShortcutUtil;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -73,20 +78,23 @@
     public DumpTarget newItemTarget(ItemInfo info) {
         DumpTarget dt = new DumpTarget();
         dt.type = DumpTarget.Type.ITEM;
-
+        if (info == null) {
+            return dt;
+        }
         switch (info.itemType) {
             case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
                 dt.itemType = ItemType.APP_ICON;
                 break;
-            case LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT:
-                dt.itemType = ItemType.UNKNOWN_ITEMTYPE;
-                break;
             case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
                 dt.itemType = ItemType.WIDGET;
                 break;
-            case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT:
+            case ITEM_TYPE_DEEP_SHORTCUT:
+            case LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT:
                 dt.itemType = ItemType.SHORTCUT;
                 break;
+            default:
+                dt.itemType = ItemType.UNKNOWN_ITEMTYPE;
+                break;
         }
         return dt;
     }
@@ -120,6 +128,9 @@
     }
 
     private static String getItemStr(DumpTarget t) {
+        if (t == null) {
+            return "";
+        }
         String typeStr = LoggerUtils.getFieldName(t.itemType, ItemType.class);
         if (!TextUtils.isEmpty(t.packageName)) {
             typeStr += ", package=" + t.packageName;
@@ -132,8 +143,15 @@
     }
 
     public DumpTarget writeToDumpTarget(ItemInfo info) {
-        node.component = info.getTargetComponent() == null? "":
-                info.getTargetComponent().flattenToString();
+        if (info == null) {
+            return node;
+        }
+        if (ShortcutUtil.isDeepShortcut(info)) {
+            node.component = ((WorkspaceItemInfo) info).getDeepShortcutId();
+        } else {
+            ComponentName cmp = info.getTargetComponent();
+            node.component = cmp == null ? "" : cmp.flattenToString();
+        }
         node.packageName = info.getTargetComponent() == null? "":
                 info.getTargetComponent().getPackageName();
         if (info instanceof LauncherAppWidgetInfo) {
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 4243ed0..60dad12 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -92,7 +92,7 @@
     public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 5;
 
-    public static final long DEFAULT_UI_TIMEOUT = 60000; // b/136278866
+    public static final long DEFAULT_UI_TIMEOUT = 10000;
     private static final String TAG = "AbstractLauncherUiTest";
 
     protected LooperExecutor mMainThreadExecutor = MAIN_EXECUTOR;
@@ -259,7 +259,7 @@
     protected <T> T getOnUiThread(final Callable<T> callback) {
         try {
             return mMainThreadExecutor.submit(callback).get();
-        } catch (Exception e) {
+        } catch (Throwable e) {
             throw new RuntimeException(e);
         }
     }
diff --git a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
index 80bb3ed..1a68122 100644
--- a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
+++ b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
@@ -38,8 +38,8 @@
 
                     evaluateInPortrait();
                     evaluateInLandscape();
-                } catch (Exception e) {
-                    Log.e(TAG, "Exception", e);
+                } catch (Throwable e) {
+                    Log.e(TAG, "Error", e);
                     throw e;
                 } finally {
                     mTest.mDevice.setOrientationNatural();
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 0472ce1..62e2a53 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -103,11 +103,11 @@
 
         setResult(acceptConfig);
         if (acceptConfig) {
-            Wait.atMost(null, new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+            Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
             assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
         } else {
             // Verify that the widget id is deleted.
-            Wait.atMost(null, () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
+            Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
                     DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
         }
     }
diff --git a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
index d909ad7..59b861c 100644
--- a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
@@ -170,7 +170,7 @@
 
         // Go back to home
         mLauncher.pressHome();
-        Wait.atMost(null, new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT,
+        Wait.atMost("", new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT,
                 mLauncher);
     }
 
diff --git a/tests/src/com/android/launcher3/util/Wait.java b/tests/src/com/android/launcher3/util/Wait.java
index 2663d02..2ab1e00 100644
--- a/tests/src/com/android/launcher3/util/Wait.java
+++ b/tests/src/com/android/launcher3/util/Wait.java
@@ -7,6 +7,8 @@
 
 import org.junit.Assert;
 
+import java.util.function.Supplier;
+
 /**
  * A utility class for waiting for a condition to be true.
  */
@@ -16,10 +18,16 @@
 
     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(String message, Condition condition, long timeout, long sleepMillis,
+    public static void atMost(Supplier<String> message, Condition condition, long timeout,
+            long sleepMillis,
             LauncherInstrumentation launcher) {
         final long startTime = SystemClock.uptimeMillis();
         long endTime = startTime + timeout;
@@ -45,6 +53,6 @@
         }
         Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis());
         launcher.checkForAnomaly();
-        Assert.fail(message);
+        Assert.fail(message.get());
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 9c550cc..b715de0 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -308,13 +308,22 @@
         }
     }
 
+    private String getVisiblePackages() {
+        return mDevice.findObjects(By.textStartsWith(""))
+                .stream()
+                .map(object -> object.getApplicationPackage())
+                .distinct()
+                .filter(pkg -> !"com.android.systemui".equals(pkg))
+                .collect(Collectors.joining(", "));
+    }
+
     private String getVisibleStateMessage() {
         if (hasLauncherObject(CONTEXT_MENU_RES_ID)) return "Context Menu";
         if (hasLauncherObject(WIDGETS_RES_ID)) return "Widgets";
         if (hasLauncherObject(OVERVIEW_RES_ID)) return "Overview";
         if (hasLauncherObject(WORKSPACE_RES_ID)) return "Workspace";
         if (hasLauncherObject(APPS_RES_ID)) return "AllApps";
-        return "Background";
+        return "Background (" + getVisiblePackages() + ")";
     }
 
     public void setSystemHealthSupplier(Function<Long, String> supplier) {
@@ -429,12 +438,6 @@
         assertEquals("Unexpected display rotation",
                 mExpectedRotation, mDevice.getDisplayRotation());
 
-        // b/136278866
-        for (int i = 0; i != 100; ++i) {
-            if (getNavigationModeMismatchError() == null) break;
-            sleep(100);
-        }
-
         final String error = getNavigationModeMismatchError();
         assertTrue(error, error == null);
         log("verifyContainerType: " + containerType);
@@ -559,7 +562,7 @@
             if (hasLauncherObject(WORKSPACE_RES_ID)) {
                 log(action = "already at home");
             } else {
-                log("Hierarchy before swiping up to home");
+                log("Hierarchy before swiping up to home:");
                 dumpViewHierarchy();
                 log(action = "swiping up to home from " + getVisibleStateMessage());
 
@@ -571,15 +574,19 @@
                 }
             }
         } else {
-            log(action = "clicking home button");
-            executeAndWaitForEvent(
-                    () -> {
-                        log("LauncherInstrumentation.pressHome before clicking");
-                        waitForSystemUiObject("home").click();
-                    },
-                    event -> true,
-                    () -> "Pressing Home didn't produce any events");
-            mDevice.waitForIdle();
+            log("Hierarchy before clicking home:");
+            dumpViewHierarchy();
+            log(action = "clicking home button from " + getVisibleStateMessage());
+            try (LauncherInstrumentation.Closable c = addContextLayer(action)) {
+                mDevice.waitForIdle();
+                runToState(
+                        () -> waitForSystemUiObject("home").click(),
+                        NORMAL_STATE_ORDINAL,
+                        !hasLauncherObject(WORKSPACE_RES_ID)
+                                && (hasLauncherObject(APPS_RES_ID)
+                                || hasLauncherObject(OVERVIEW_RES_ID)));
+                mDevice.waitForIdle();
+            }
         }
         try (LauncherInstrumentation.Closable c = addContextLayer(
                 "performed action to switch to Home - " + action)) {
@@ -783,12 +790,20 @@
                 + "]";
     }
 
+    void runToState(Runnable command, int expectedState, boolean requireEvent) {
+        if (requireEvent) {
+            runToState(command, expectedState);
+        } else {
+            command.run();
+        }
+    }
+
     void runToState(Runnable command, int expectedState) {
         final List<Integer> actualEvents = new ArrayList<>();
         executeAndWaitForEvent(
                 command,
                 event -> isSwitchToStateEvent(event, expectedState, actualEvents),
-                () -> "Failed to receive an event for the swipe end: expected "
+                () -> "Failed to receive an event for the state change: expected "
                         + TestProtocol.stateOrdinalToString(expectedState)
                         + ", actual: " + eventListToString(actualEvents));
     }