diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java
index fa28106..af67e1b 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java
@@ -27,11 +27,13 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.os.Build;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.Message;
 import android.os.UserHandle;
 import android.util.Log;
 
+import androidx.annotation.Nullable;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.appprediction.PredictionUiStateManager.Client;
 import com.android.launcher3.model.AppLaunchTracker;
@@ -97,6 +99,7 @@
                 new AppPredictionContext.Builder(mContext)
                         .setUiSurface(client.id)
                         .setPredictedTargetCount(count)
+                        .setExtras(getAppPredictionContextExtras(client))
                         .build());
         predictor.registerPredictionUpdates(mContext.getMainExecutor(),
                 PredictionUiStateManager.INSTANCE.get(mContext).appPredictorCallback(client));
@@ -104,6 +107,15 @@
         return predictor;
     }
 
+    /**
+     * Override to add custom extras.
+     */
+    @WorkerThread
+    @Nullable
+    public Bundle getAppPredictionContextExtras(Client client){
+        return null;
+    }
+
     @WorkerThread
     private boolean handleMessage(Message msg) {
         switch (msg.what) {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
index 537858d..29034ea 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -176,7 +176,8 @@
         RECENTS(1, STATE_SCALED_CONTROLLER_RECENTS | STATE_CAPTURE_SCREENSHOT
                 | STATE_SCREENSHOT_VIEW_SHOWN, true, false, ContainerType.TASKSWITCHER, true),
 
-        NEW_TASK(0, STATE_START_NEW_TASK, false, true, ContainerType.APP, true),
+        NEW_TASK(0, STATE_START_NEW_TASK | STATE_CAPTURE_SCREENSHOT, false, true,
+                ContainerType.APP, true),
 
         LAST_TASK(0, STATE_RESUME_LAST_TASK, false, true, ContainerType.APP, false);
 
@@ -320,7 +321,7 @@
 
         mStateCallback.addCallback(STATE_RESUME_LAST_TASK | STATE_APP_CONTROLLER_RECEIVED,
                 this::resumeLastTask);
-        mStateCallback.addCallback(STATE_START_NEW_TASK | STATE_APP_CONTROLLER_RECEIVED,
+        mStateCallback.addCallback(STATE_START_NEW_TASK | STATE_SCREENSHOT_CAPTURED,
                 this::startNewTask);
 
         mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_APP_CONTROLLER_RECEIVED
diff --git a/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java b/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java
index 47ce44c..5e20e56 100644
--- a/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java
+++ b/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java
@@ -79,6 +79,8 @@
 
     @After
     public void tearDown() throws Throwable {
+        AppLaunchTracker.INSTANCE.initializeForTesting(null);
+        PredictionUiStateManager.INSTANCE.initializeForTesting(null);
         mDevice.unfreezeRotation();
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index ec62188..d3f1460 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -58,6 +58,7 @@
 
 import org.junit.Assert;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
 import java.util.Deque;
@@ -214,7 +215,22 @@
         };
     }
 
+    private void dumpViewHierarchy() {
+        final ByteArrayOutputStream stream = new ByteArrayOutputStream();
+        try {
+            mDevice.dumpWindowHierarchy(stream);
+            stream.flush();
+            stream.close();
+            for (String line : stream.toString().split("\\r?\\n")) {
+                Log.e(TAG, line.trim());
+            }
+        } catch (IOException e) {
+            Log.e(TAG, "error dumping XML to logcat", e);
+        }
+    }
+
     private void fail(String message) {
+        dumpViewHierarchy();
         Assert.fail("http://go/tapl : " + getContextDescription() + message);
     }