Merge "Reducing chances for the object with finalizer to be reffed by register" into ub-launcher3-master
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
index 1d00825..c94c6d5 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
@@ -25,7 +25,6 @@
 
 import android.animation.Animator;
 import android.annotation.TargetApi;
-import android.app.ActivityManager.RunningTaskInfo;
 import android.content.Context;
 import android.content.Intent;
 import android.graphics.Point;
@@ -96,7 +95,6 @@
     protected final OverviewComponentObserver mOverviewComponentObserver;
     protected final BaseActivityInterface<T> mActivityInterface;
     protected final RecentsModel mRecentsModel;
-    protected final int mRunningTaskId;
 
     protected final AppWindowAnimationHelper mAppWindowAnimationHelper;
     protected final TransformParams mTransformParams = new TransformParams();
@@ -130,7 +128,7 @@
 
     protected BaseSwipeUpHandler(Context context, RecentsAnimationDeviceState deviceState,
             GestureState gestureState, OverviewComponentObserver overviewComponentObserver,
-            RecentsModel recentsModel, InputConsumerController inputConsumer, int runningTaskId) {
+            RecentsModel recentsModel, InputConsumerController inputConsumer) {
         mContext = context;
         mDeviceState = deviceState;
         mGestureState = gestureState;
@@ -139,7 +137,6 @@
         mRecentsModel = recentsModel;
         mActivityInitListener =
                 mActivityInterface.createActivityInitListener(this::onActivityInit);
-        mRunningTaskId = runningTaskId;
         mInputConsumer = inputConsumer;
 
         mAppWindowAnimationHelper = new AppWindowAnimationHelper(context);
@@ -261,7 +258,8 @@
         mRecentsAnimationTargets = targets;
         DeviceProfile dp = InvariantDeviceProfile.INSTANCE.get(mContext).getDeviceProfile(mContext);
         final Rect overviewStackBounds;
-        RemoteAnimationTargetCompat runningTaskTarget = targets.findTask(mRunningTaskId);
+        RemoteAnimationTargetCompat runningTaskTarget = targets.findTask(
+                mGestureState.getRunningTaskId());
 
         if (targets.minimizedHomeBounds != null && runningTaskTarget != null) {
             overviewStackBounds = mActivityInterface
@@ -490,8 +488,8 @@
 
     public interface Factory {
 
-        BaseSwipeUpHandler newHandler(GestureState gestureState, RunningTaskInfo runningTask,
-                long touchTimeMs, boolean continuingLastGesture, boolean isLikelyToStartNewTask);
+        BaseSwipeUpHandler newHandler(GestureState gestureState, long touchTimeMs,
+                boolean continuingLastGesture, boolean isLikelyToStartNewTask);
     }
 
     protected interface RunningWindowAnim {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
index b2626e5..a97b0af 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
@@ -434,6 +434,8 @@
             GestureState newGestureState = new GestureState(
                     mOverviewComponentObserver.getActivityInterface(),
                     ActiveGestureLog.INSTANCE.generateAndSetLogId());
+            newGestureState.updateRunningTask(TraceHelper.whitelistIpcs("getRunningTask.0",
+                    () -> mAM.getRunningTask(0)));
 
             if (mDeviceState.isInSwipeUpTouchRegion(event)) {
                 mConsumer.onConsumerAboutToBeSwitched();
@@ -470,8 +472,7 @@
             if (canStartSystemGesture) {
                 // This handles apps launched in direct boot mode (e.g. dialer) as well as apps
                 // launched while device is locked even after exiting direct boot mode (e.g. camera).
-                return createDeviceLockedInputConsumer(newGestureState,
-                        mAM.getRunningTask(ACTIVITY_TYPE_ASSISTANT));
+                return createDeviceLockedInputConsumer(newGestureState);
             } else {
                 return mResetGestureInputConsumer;
             }
@@ -514,25 +515,22 @@
 
     private InputConsumer newBaseConsumer(GestureState previousGestureState,
             GestureState gestureState, MotionEvent event) {
-        RunningTaskInfo runningTaskInfo = TraceHelper.whitelistIpcs("getRunningTask.0",
-                () -> mAM.getRunningTask(0));
         if (mDeviceState.isKeyguardShowingOccluded()) {
             // This handles apps showing over the lockscreen (e.g. camera)
-            return createDeviceLockedInputConsumer(gestureState, runningTaskInfo);
+            return createDeviceLockedInputConsumer(gestureState);
         }
 
         boolean forceOverviewInputConsumer = false;
-        if (isExcludedAssistant(runningTaskInfo)) {
+        if (isExcludedAssistant(gestureState.getRunningTask())) {
             // In the case where we are in the excluded assistant state, ignore it and treat the
             // running activity as the task behind the assistant
-
-            runningTaskInfo = TraceHelper.whitelistIpcs("getRunningTask.assistant",
-                    () -> mAM.getRunningTask(ACTIVITY_TYPE_ASSISTANT));
-            if (!ActivityManagerWrapper.isHomeTask(runningTaskInfo)) {
+            gestureState.updateRunningTask(TraceHelper.whitelistIpcs("getRunningTask.assistant",
+                    () -> mAM.getRunningTask(ACTIVITY_TYPE_ASSISTANT)));
+            if (!ActivityManagerWrapper.isHomeTask(gestureState.getRunningTask())) {
                 final ComponentName homeComponent =
                         mOverviewComponentObserver.getHomeIntent().getComponent();
-                forceOverviewInputConsumer =
-                        runningTaskInfo.baseIntent.getComponent().equals(homeComponent);
+                forceOverviewInputConsumer = gestureState.getRunningTask()
+                        .baseIntent.getComponent().equals(homeComponent);
             }
         }
 
@@ -541,9 +539,9 @@
             // consumer but with the next task as the running task
             RunningTaskInfo info = new ActivityManager.RunningTaskInfo();
             info.id = previousGestureState.getFinishingRecentsAnimationTaskId();
-            return createOtherActivityInputConsumer(previousGestureState, gestureState, event,
-                    info);
-        } else if (runningTaskInfo == null) {
+            gestureState.updateRunningTask(info);
+            return createOtherActivityInputConsumer(previousGestureState, gestureState, event);
+        } else if (gestureState.getRunningTask() == null) {
             return mResetGestureInputConsumer;
         } else if (previousGestureState.isRunningAnimationToLauncher()
                 || gestureState.getActivityInterface().isResumed()
@@ -552,11 +550,10 @@
         } else if (ENABLE_QUICKSTEP_LIVE_TILE.get()
                 && gestureState.getActivityInterface().isInLiveTileMode()) {
             return createOverviewInputConsumer(previousGestureState, gestureState, event);
-        } else if (mDeviceState.isGestureBlockedActivity(runningTaskInfo)) {
+        } else if (mDeviceState.isGestureBlockedActivity(gestureState.getRunningTask())) {
             return mResetGestureInputConsumer;
         } else {
-            return createOtherActivityInputConsumer(previousGestureState, gestureState, event,
-                    runningTaskInfo);
+            return createOtherActivityInputConsumer(previousGestureState, gestureState, event);
         }
     }
 
@@ -567,8 +564,7 @@
     }
 
     private InputConsumer createOtherActivityInputConsumer(GestureState previousGestureState,
-            GestureState gestureState,
-            MotionEvent event, RunningTaskInfo runningTaskInfo) {
+            GestureState gestureState, MotionEvent event) {
 
         final boolean shouldDefer;
         final BaseSwipeUpHandler.Factory factory;
@@ -585,15 +581,14 @@
 
         final boolean disableHorizontalSwipe = mDeviceState.isInExclusionRegion(event);
         return new OtherActivityInputConsumer(this, mDeviceState, mTaskAnimationManager,
-                gestureState, runningTaskInfo, shouldDefer, this::onConsumerInactive,
+                gestureState, shouldDefer, this::onConsumerInactive,
                 mInputMonitorCompat, disableHorizontalSwipe, factory);
     }
 
-    private InputConsumer createDeviceLockedInputConsumer(GestureState gestureState,
-            RunningTaskInfo taskInfo) {
-        if (mDeviceState.isFullyGesturalNavMode() && taskInfo != null) {
+    private InputConsumer createDeviceLockedInputConsumer(GestureState gestureState) {
+        if (mDeviceState.isFullyGesturalNavMode() && gestureState.getRunningTask() != null) {
             return new DeviceLockedInputConsumer(this, mDeviceState, mTaskAnimationManager,
-                    gestureState, mInputMonitorCompat, taskInfo.taskId);
+                    gestureState, mInputMonitorCompat);
         } else {
             return mResetGestureInputConsumer;
         }
@@ -727,19 +722,17 @@
     }
 
     private BaseSwipeUpHandler createWindowTransformSwipeHandler(GestureState gestureState,
-            RunningTaskInfo runningTask, long touchTimeMs, boolean continuingLastGesture,
-            boolean isLikelyToStartNewTask) {
+            long touchTimeMs, boolean continuingLastGesture, boolean isLikelyToStartNewTask) {
         return  new WindowTransformSwipeHandler(this, mDeviceState, mTaskAnimationManager,
-                gestureState, runningTask, touchTimeMs, mOverviewComponentObserver,
-                continuingLastGesture, mInputConsumer, mRecentsModel);
+                gestureState, touchTimeMs, mOverviewComponentObserver, continuingLastGesture,
+                mInputConsumer, mRecentsModel);
     }
 
     private BaseSwipeUpHandler createFallbackNoButtonSwipeHandler(GestureState gestureState,
-            RunningTaskInfo runningTask, long touchTimeMs, boolean continuingLastGesture,
-            boolean isLikelyToStartNewTask) {
+            long touchTimeMs, boolean continuingLastGesture, boolean isLikelyToStartNewTask) {
         return new FallbackNoButtonInputConsumer(this, mDeviceState, gestureState,
-                mOverviewComponentObserver, runningTask, mRecentsModel, mInputConsumer,
-                isLikelyToStartNewTask, continuingLastGesture);
+                mOverviewComponentObserver, mRecentsModel, mInputConsumer, isLikelyToStartNewTask,
+                continuingLastGesture);
     }
 
     protected boolean shouldNotifyBackGesture() {
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 065f3eb..1e636e1 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -39,7 +39,6 @@
 import android.animation.TimeInterpolator;
 import android.animation.ValueAnimator;
 import android.annotation.TargetApi;
-import android.app.ActivityManager.RunningTaskInfo;
 import android.content.Context;
 import android.content.Intent;
 import android.graphics.PointF;
@@ -192,11 +191,11 @@
 
     public WindowTransformSwipeHandler(Context context, RecentsAnimationDeviceState deviceState,
             TaskAnimationManager taskAnimationManager, GestureState gestureState,
-            RunningTaskInfo runningTaskInfo, long touchTimeMs,
-            OverviewComponentObserver overviewComponentObserver, boolean continuingLastGesture,
-            InputConsumerController inputConsumer, RecentsModel recentsModel) {
+            long touchTimeMs, OverviewComponentObserver overviewComponentObserver,
+            boolean continuingLastGesture, InputConsumerController inputConsumer,
+            RecentsModel recentsModel) {
         super(context, deviceState, gestureState, overviewComponentObserver, recentsModel,
-                inputConsumer, runningTaskInfo.id);
+                inputConsumer);
         mTaskAnimationManager = taskAnimationManager;
         mTouchTimeMs = touchTimeMs;
         mContinuingLastGesture = continuingLastGesture;
@@ -290,9 +289,9 @@
 
         mStateCallback.setState(STATE_LAUNCHER_PRESENT);
         if (alreadyOnHome) {
-            onLauncherStart(activity);
+            onLauncherStart();
         } else {
-            activity.setOnStartCallback(this::onLauncherStart);
+            activity.runOnceOnStart(this::onLauncherStart);
         }
 
         setupRecentsViewUi();
@@ -304,7 +303,8 @@
         return mGestureState.getEndTarget() != HOME;
     }
 
-    private void onLauncherStart(final T activity) {
+    private void onLauncherStart() {
+        final T activity = mActivityInterface.getCreatedActivity();
         if (mActivity != activity) {
             return;
         }
@@ -393,7 +393,7 @@
             updateSysUiFlags(mCurrentShift.value);
             return;
         }
-        mRecentsView.onGestureAnimationStart(mRunningTaskId);
+        mRecentsView.onGestureAnimationStart(mGestureState.getRunningTaskId());
     }
 
     private void launcherFrameDrawn() {
@@ -442,9 +442,9 @@
         if (!mDeviceState.isFullyGesturalNavMode() || mRecentsView == null) {
             return;
         }
-        RemoteAnimationTargetCompat runningTaskTarget = mRecentsAnimationTargets == null
-                ? null
-                : mRecentsAnimationTargets.findTask(mRunningTaskId);
+        RemoteAnimationTargetCompat runningTaskTarget = mRecentsAnimationTargets != null
+                ? mRecentsAnimationTargets.findTask(mGestureState.getRunningTaskId())
+                : null;
         final boolean recentsAttachedToAppWindow;
         if (mGestureState.getEndTarget() != null) {
             recentsAttachedToAppWindow = mGestureState.getEndTarget().recentsAttachedToAppWindow;
@@ -1005,7 +1005,9 @@
     @Override
     public void onConsumerAboutToBeSwitched() {
         if (mActivity != null) {
-            mActivity.setOnStartCallback(null);
+            // In the off chance that the gesture ends before Launcher is started, we should clear
+            // the callback here so that it doesn't update with the wrong state
+            mActivity.clearRunOnceOnStartCallback();
         }
         if (mGestureState.getEndTarget() != null && !mGestureState.isRunningAnimationToLauncher()) {
             cancelCurrentAnimation();
@@ -1114,13 +1116,14 @@
     }
 
     private void switchToScreenshot() {
+        final int runningTaskId = mGestureState.getRunningTaskId();
         if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
             if (mRecentsAnimationController != null) {
                 // Update the screenshot of the task
                 if (mTaskSnapshot == null) {
-                    mTaskSnapshot = mRecentsAnimationController.screenshotTask(mRunningTaskId);
+                    mTaskSnapshot = mRecentsAnimationController.screenshotTask(runningTaskId);
                 }
-                mRecentsView.updateThumbnail(mRunningTaskId, mTaskSnapshot, false /* refreshNow */);
+                mRecentsView.updateThumbnail(runningTaskId, mTaskSnapshot, false /* refreshNow */);
             }
             mStateCallback.setStateOnUiThread(STATE_SCREENSHOT_CAPTURED);
         } else if (!hasTargets()) {
@@ -1131,7 +1134,7 @@
             if (mRecentsAnimationController != null) {
                 // Update the screenshot of the task
                 if (mTaskSnapshot == null) {
-                    mTaskSnapshot = mRecentsAnimationController.screenshotTask(mRunningTaskId);
+                    mTaskSnapshot = mRecentsAnimationController.screenshotTask(runningTaskId);
                 }
                 final TaskView taskView;
                 if (mGestureState.getEndTarget() == HOME) {
@@ -1139,7 +1142,7 @@
                     // taken in the correct orientation, but no need to update the thumbnail.
                     taskView = null;
                 } else {
-                    taskView = mRecentsView.updateThumbnail(mRunningTaskId, mTaskSnapshot);
+                    taskView = mRecentsView.updateThumbnail(runningTaskId, mTaskSnapshot);
                 }
                 if (taskView != null && !mCanceled) {
                     // Defer finishing the animation until the next launcher frame with the
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
index 370f161..e448c5b 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
@@ -85,7 +85,6 @@
     private final AppWindowAnimationHelper.TransformParams mTransformParams;
     private final Point mDisplaySize;
     private final MultiStateCallback mStateCallback;
-    public final int mRunningTaskId;
 
     private VelocityTracker mVelocityTracker;
     private float mProgress;
@@ -97,7 +96,7 @@
 
     public DeviceLockedInputConsumer(Context context, RecentsAnimationDeviceState deviceState,
             TaskAnimationManager taskAnimationManager, GestureState gestureState,
-            InputMonitorCompat inputMonitorCompat, int runningTaskId) {
+            InputMonitorCompat inputMonitorCompat) {
         mContext = context;
         mDeviceState = deviceState;
         mTaskAnimationManager = taskAnimationManager;
@@ -106,7 +105,6 @@
         mAppWindowAnimationHelper = new AppWindowAnimationHelper(context);
         mTransformParams = new AppWindowAnimationHelper.TransformParams();
         mInputMonitorCompat = inputMonitorCompat;
-        mRunningTaskId = runningTaskId;
 
         // Do not use DeviceProfile as the user data might be locked
         mDisplaySize = DefaultDisplay.INSTANCE.get(context).getInfo().realSize;
@@ -221,7 +219,8 @@
         mRecentsAnimationTargets = targets;
 
         Rect displaySize = new Rect(0, 0, mDisplaySize.x, mDisplaySize.y);
-        RemoteAnimationTargetCompat targetCompat = targets.findTask(mRunningTaskId);
+        RemoteAnimationTargetCompat targetCompat = targets.findTask(
+                mGestureState.getRunningTaskId());
         if (targetCompat != null) {
             mAppWindowAnimationHelper.updateSource(displaySize, targetCompat);
         }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java
index 7b24bd9..5ffb975 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java
@@ -27,7 +27,6 @@
 
 import android.animation.Animator;
 import android.animation.AnimatorSet;
-import android.app.ActivityManager.RunningTaskInfo;
 import android.app.ActivityOptions;
 import android.content.Context;
 import android.content.Intent;
@@ -108,24 +107,20 @@
     private final boolean mRunningOverHome;
     private final boolean mSwipeUpOverHome;
 
-    private final RunningTaskInfo mRunningTaskInfo;
-
     private final PointF mEndVelocityPxPerMs = new PointF(0, 0.5f);
     private RunningWindowAnim mFinishAnimation;
 
     public FallbackNoButtonInputConsumer(Context context, RecentsAnimationDeviceState deviceState,
             GestureState gestureState, OverviewComponentObserver overviewComponentObserver,
-            RunningTaskInfo runningTaskInfo, RecentsModel recentsModel,
-            InputConsumerController inputConsumer,
+            RecentsModel recentsModel, InputConsumerController inputConsumer,
             boolean isLikelyToStartNewTask, boolean continuingLastGesture) {
         super(context, deviceState, gestureState, overviewComponentObserver, recentsModel,
-                inputConsumer, runningTaskInfo.id);
+                inputConsumer);
         mLauncherAlpha.value = 1;
 
-        mRunningTaskInfo = runningTaskInfo;
         mInQuickSwitchMode = isLikelyToStartNewTask || continuingLastGesture;
         mContinuingLastGesture = continuingLastGesture;
-        mRunningOverHome = ActivityManagerWrapper.isHomeTask(runningTaskInfo);
+        mRunningOverHome = ActivityManagerWrapper.isHomeTask(mGestureState.getRunningTask());
         mSwipeUpOverHome = mRunningOverHome && !mInQuickSwitchMode;
 
         if (mSwipeUpOverHome) {
@@ -182,9 +177,9 @@
 
         if (!mContinuingLastGesture) {
             if (mRunningOverHome) {
-                mRecentsView.onGestureAnimationStart(mRunningTaskInfo);
+                mRecentsView.onGestureAnimationStart(mGestureState.getRunningTask());
             } else {
-                mRecentsView.onGestureAnimationStart(mRunningTaskId);
+                mRecentsView.onGestureAnimationStart(mGestureState.getRunningTaskId());
             }
         }
         mStateCallback.setStateOnUiThread(STATE_RECENTS_PRESENT);
@@ -344,7 +339,8 @@
                     break;
                 }
 
-                ThumbnailData thumbnail = mRecentsAnimationController.screenshotTask(mRunningTaskId);
+                final int runningTaskId = mGestureState.getRunningTaskId();
+                ThumbnailData thumbnail = mRecentsAnimationController.screenshotTask(runningTaskId);
                 mRecentsAnimationController.setDeferCancelUntilNextTransition(true /* defer */,
                         false /* screenshot */);
 
@@ -353,7 +349,7 @@
 
                 Bundle extras = new Bundle();
                 extras.putBinder(EXTRA_THUMBNAIL, new ObjectWrapper<>(thumbnail));
-                extras.putInt(EXTRA_TASK_ID, mRunningTaskId);
+                extras.putInt(EXTRA_TASK_ID, runningTaskId);
 
                 Intent intent = new Intent(mOverviewComponentObserver.getOverviewIntent())
                         .putExtras(extras);
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index aeab4b5..f8676ec 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -30,7 +30,6 @@
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
 
 import android.annotation.TargetApi;
-import android.app.ActivityManager.RunningTaskInfo;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.Intent;
@@ -81,7 +80,6 @@
     private final GestureState mGestureState;
     private RecentsAnimationCallbacks mActiveCallbacks;
     private final CachedEventDispatcher mRecentsViewDispatcher = new CachedEventDispatcher();
-    private final RunningTaskInfo mRunningTask;
     private final InputMonitorCompat mInputMonitorCompat;
     private final BaseActivityInterface mActivityInterface;
 
@@ -124,8 +122,7 @@
 
     public OtherActivityInputConsumer(Context base, RecentsAnimationDeviceState deviceState,
             TaskAnimationManager taskAnimationManager, GestureState gestureState,
-            RunningTaskInfo runningTaskInfo, boolean isDeferredDownTarget,
-            Consumer<OtherActivityInputConsumer> onCompleteCallback,
+            boolean isDeferredDownTarget, Consumer<OtherActivityInputConsumer> onCompleteCallback,
             InputMonitorCompat inputMonitorCompat, boolean disableHorizontalSwipe,
             Factory handlerFactory) {
         super(base);
@@ -133,7 +130,6 @@
         mTaskAnimationManager = taskAnimationManager;
         mGestureState = gestureState;
         mMainThreadHandler = new Handler(Looper.getMainLooper());
-        mRunningTask = runningTaskInfo;
         mHandlerFactory = handlerFactory;
         mActivityInterface = mGestureState.getActivityInterface();
 
@@ -325,7 +321,7 @@
             long touchTimeMs, boolean isLikelyToStartNewTask) {
         ActiveGestureLog.INSTANCE.addLog("startRecentsAnimation");
 
-        mInteractionHandler = mHandlerFactory.newHandler(mGestureState, mRunningTask, touchTimeMs,
+        mInteractionHandler = mHandlerFactory.newHandler(mGestureState, touchTimeMs,
                 mTaskAnimationManager.isRecentsAnimationRunning(), isLikelyToStartNewTask);
         mInteractionHandler.setGestureEndCallback(this::onInteractionGestureFinished);
         mMotionPauseDetector.setOnMotionPauseListener(mInteractionHandler::onMotionPauseChanged);
diff --git a/quickstep/src/com/android/quickstep/GestureState.java b/quickstep/src/com/android/quickstep/GestureState.java
index 98ff410..28ed32c 100644
--- a/quickstep/src/com/android/quickstep/GestureState.java
+++ b/quickstep/src/com/android/quickstep/GestureState.java
@@ -17,6 +17,7 @@
 
 import static com.android.quickstep.MultiStateCallback.DEBUG_STATES;
 
+import android.app.ActivityManager;
 import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.userevent.nano.LauncherLogProto.ContainerType;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -102,21 +103,22 @@
     private final MultiStateCallback mStateCallback;
     private final int mGestureId;
 
+    private ActivityManager.RunningTaskInfo mRunningTask;
     private GestureEndTarget mEndTarget;
     // TODO: This can be removed once we stop finishing the animation when starting a new task
     private int mFinishingRecentsAnimationTaskId = -1;
 
     public GestureState(BaseActivityInterface activityInterface, int gestureId) {
         mActivityInterface = activityInterface;
-        mGestureId = gestureId;
         mStateCallback = new MultiStateCallback(STATE_NAMES.toArray(new String[0]));
+        mGestureId = gestureId;
     }
 
     public GestureState() {
         // Do nothing, only used for initializing the gesture state prior to user unlock
         mActivityInterface = null;
-        mGestureId = -1;
         mStateCallback = new MultiStateCallback(STATE_NAMES.toArray(new String[0]));
+        mGestureId = -1;
     }
 
     /**
@@ -148,6 +150,27 @@
     }
 
     /**
+     * @return the running task for this gesture.
+     */
+    public ActivityManager.RunningTaskInfo getRunningTask() {
+        return mRunningTask;
+    }
+
+    /**
+     * @return the running task id for this gesture.
+     */
+    public int getRunningTaskId() {
+        return mRunningTask != null ? mRunningTask.taskId : -1;
+    }
+
+    /**
+     * Updates the running task for the gesture to be the given {@param runningTask}.
+     */
+    public void updateRunningTask(ActivityManager.RunningTaskInfo runningTask) {
+        mRunningTask = runningTask;
+    }
+
+    /**
      * @return the end target for this gesture (if known).
      */
     public GestureEndTarget getEndTarget() {
@@ -155,14 +178,6 @@
     }
 
     /**
-     * @return whether the current gesture is still running a recents animation to a state in the
-     *         Launcher or Recents activity.
-     */
-    public boolean isRunningAnimationToLauncher() {
-        return isRecentsAnimationRunning() && mEndTarget != null && mEndTarget.isLauncher;
-    }
-
-    /**
      * Sets the end target of this gesture and immediately notifies the state changes.
      */
     public void setEndTarget(GestureEndTarget target) {
@@ -202,6 +217,15 @@
     }
 
     /**
+     * @return whether the current gesture is still running a recents animation to a state in the
+     *         Launcher or Recents activity.
+     * Updates the running task for the gesture to be the given {@param runningTask}.
+     */
+    public boolean isRunningAnimationToLauncher() {
+        return isRecentsAnimationRunning() && mEndTarget != null && mEndTarget.isLauncher;
+    }
+
+    /**
      * @return whether the recents animation is started but not yet ended
      */
     public boolean isRecentsAnimationRunning() {
diff --git a/src/com/android/launcher3/BaseDraggingActivity.java b/src/com/android/launcher3/BaseDraggingActivity.java
index d24de8e..772eb00 100644
--- a/src/com/android/launcher3/BaseDraggingActivity.java
+++ b/src/com/android/launcher3/BaseDraggingActivity.java
@@ -57,7 +57,7 @@
     private ActionMode mCurrentActionMode;
     protected boolean mIsSafeModeEnabled;
 
-    private OnStartCallback mOnStartCallback;
+    private Runnable mOnStartCallback;
 
     private int mThemeRes = R.style.AppTheme;
 
@@ -226,7 +226,7 @@
         super.onStart();
 
         if (mOnStartCallback != null) {
-            mOnStartCallback.onActivityStart(this);
+            mOnStartCallback.run();
             mOnStartCallback = null;
         }
     }
@@ -238,8 +238,12 @@
         mRotationListener.disable();
     }
 
-    public <T extends BaseDraggingActivity> void setOnStartCallback(OnStartCallback<T> callback) {
-        mOnStartCallback = callback;
+    public void runOnceOnStart(Runnable action) {
+        mOnStartCallback = action;
+    }
+
+    public void clearRunOnceOnStartCallback() {
+        mOnStartCallback = null;
     }
 
     protected void onDeviceProfileInitiated() {
@@ -258,12 +262,4 @@
     }
 
     protected abstract void reapplyUi();
-
-    /**
-     * Callback for listening for onStart
-     */
-    public interface OnStartCallback<T extends BaseDraggingActivity> {
-
-        void onActivityStart(T activity);
-    }
 }
diff --git a/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java b/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java
index 08fa098..858cb38 100644
--- a/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java
+++ b/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java
@@ -21,6 +21,7 @@
 import android.os.Build;
 import android.util.Log;
 
+import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.UiDevice;
 
 import org.junit.rules.TestRule;
@@ -86,6 +87,19 @@
     }
 
     private static int getRunFlavor() {
+        final String flavorOverride = InstrumentationRegistry.getArguments().getString("flavor");
+
+        if (flavorOverride != null) {
+            Log.d(TAG, "Flavor override: " + flavorOverride);
+            try {
+                return (int) TestStabilityRule.class.getField(flavorOverride).get(null);
+            } catch (NoSuchFieldException e) {
+                throw new AssertionError("Unrecognized run flavor override: " + flavorOverride);
+            } catch (IllegalAccessException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
         final String launcherVersion;
         try {
             launcherVersion = getInstrumentation().