diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/LauncherInitListenerEx.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/LauncherInitListenerEx.java
deleted file mode 100644
index 76050d5..0000000
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/LauncherInitListenerEx.java
+++ /dev/null
@@ -1,34 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3;
-
-import com.android.launcher3.appprediction.PredictionUiStateManager;
-import com.android.launcher3.appprediction.PredictionUiStateManager.Client;
-
-import java.util.function.BiPredicate;
-
-public class LauncherInitListenerEx extends LauncherInitListener {
-
-    public LauncherInitListenerEx(BiPredicate<Launcher, Boolean> onInitListener) {
-        super(onInitListener);
-    }
-
-    @Override
-    public boolean init(Launcher launcher, boolean alreadyOnHome) {
-        PredictionUiStateManager.INSTANCE.get(launcher).switchClient(Client.OVERVIEW);
-        return super.init(launcher, alreadyOnHome);
-    }
-}
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityInterface.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityInterface.java
index 8502c08..01cf4bb 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityInterface.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityInterface.java
@@ -49,12 +49,13 @@
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherInitListenerEx;
+import com.android.launcher3.LauncherInitListener;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.allapps.DiscoveryBounce;
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.AnimatorSetBuilder;
+import com.android.launcher3.appprediction.PredictionUiStateManager;
 import com.android.launcher3.uioverrides.states.OverviewState;
 import com.android.launcher3.userevent.nano.LauncherLogProto;
 import com.android.launcher3.views.FloatingIconView;
@@ -413,7 +414,7 @@
 
     @Override
     public ActivityInitListener createActivityInitListener(Predicate<Boolean> onInitListener) {
-        return new LauncherInitListenerEx((activity, alreadyOnHome) ->
+        return new LauncherInitListener((activity, alreadyOnHome) ->
                 onInitListener.test(alreadyOnHome));
     }
 
@@ -541,4 +542,14 @@
         }
         launcher.setOnDeferredActivityLaunchCallback(r);
     }
+
+    @Override
+    public void updateOverviewPredictionState() {
+        Launcher launcher = getCreatedActivity();
+        if (launcher == null) {
+            return;
+        }
+        PredictionUiStateManager.INSTANCE.get(launcher).switchClient(
+                PredictionUiStateManager.Client.OVERVIEW);
+    }
 }
\ No newline at end of file
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 7531741..a3db6ab 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
@@ -32,6 +32,7 @@
 import static com.android.quickstep.GestureState.GestureEndTarget.RECENTS;
 import static com.android.quickstep.GestureState.STATE_END_TARGET_ANIMATION_FINISHED;
 import static com.android.quickstep.MultiStateCallback.DEBUG_STATES;
+import static com.android.quickstep.SysUINavigationMode.Mode.TWO_BUTTONS;
 import static com.android.quickstep.views.RecentsView.UPDATE_SYSUI_FLAGS_THRESHOLD;
 
 import android.animation.Animator;
@@ -296,6 +297,12 @@
         }
 
         setupRecentsViewUi();
+
+        if (mDeviceState.getNavMode() == TWO_BUTTONS) {
+            // If the device is in two button mode, swiping up will show overview with predictions
+            // so we need to kick off switching to the overview predictions as soon as possible
+            mActivityInterface.updateOverviewPredictionState();
+        }
         return true;
     }
 
@@ -426,6 +433,12 @@
     @Override
     public void onMotionPauseChanged(boolean isPaused) {
         setShelfState(isPaused ? PEEK : HIDE, OVERSHOOT_1_2, SHELF_ANIM_DURATION);
+
+        if (mDeviceState.isFullyGesturalNavMode() && isPaused) {
+            // In fully gestural nav mode, switch to overview predictions once the user has paused
+            // (this is a no-op if the predictions are already in that state)
+            mActivityInterface.updateOverviewPredictionState();
+        }
     }
 
     public void maybeUpdateRecentsAttachedState() {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/OverviewCommandHelper.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/OverviewCommandHelper.java
index 7f53d83..c4733bd 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/OverviewCommandHelper.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/OverviewCommandHelper.java
@@ -29,6 +29,7 @@
 
 import androidx.annotation.BinderThread;
 import com.android.launcher3.BaseDraggingActivity;
+import com.android.launcher3.appprediction.PredictionUiStateManager;
 import com.android.launcher3.logging.UserEventDispatcher;
 import com.android.launcher3.userevent.nano.LauncherLogProto;
 import com.android.quickstep.util.ActivityInitListener;
@@ -212,6 +213,10 @@
                         LauncherLogProto.ContainerType.TASKSWITCHER);
                 mUserEventLogged = true;
             }
+
+            // Switch prediction client to overview
+            PredictionUiStateManager.INSTANCE.get(activity).switchClient(
+                    PredictionUiStateManager.Client.OVERVIEW);
             return mAnimationProvider.onActivityReady(activity, wasVisible);
         }
 
diff --git a/quickstep/src/com/android/launcher3/LauncherInitListener.java b/quickstep/src/com/android/launcher3/LauncherInitListener.java
index 663b125..96340b2 100644
--- a/quickstep/src/com/android/launcher3/LauncherInitListener.java
+++ b/quickstep/src/com/android/launcher3/LauncherInitListener.java
@@ -22,6 +22,7 @@
 import android.os.CancellationSignal;
 import android.os.Handler;
 
+import com.android.launcher3.util.ActivityTracker;
 import com.android.quickstep.util.ActivityInitListener;
 import com.android.quickstep.util.RemoteAnimationProvider;
 
@@ -32,6 +33,11 @@
 
     private RemoteAnimationProvider mRemoteAnimationProvider;
 
+    /**
+     * @param onInitListener a callback made when the activity is initialized. The callback should
+     *                       return true to continue receiving callbacks (ie. for if the activity is
+     *                       recreated).
+     */
     public LauncherInitListener(BiPredicate<Launcher, Boolean> onInitListener) {
         super(onInitListener, Launcher.ACTIVITY_TRACKER);
     }
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index cb18001..86ffd22 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -89,6 +89,13 @@
     }
 
     /**
+     * Updates the prediction state to the overview state.
+     */
+    default void updateOverviewPredictionState() {
+        // By default overview predictions are not supported
+    }
+
+    /**
      * Used for containerType in {@link com.android.launcher3.logging.UserEventDispatcher}
      */
     int getContainerType();
diff --git a/quickstep/src/com/android/quickstep/util/ActivityInitListener.java b/quickstep/src/com/android/quickstep/util/ActivityInitListener.java
index fe37d60..b1c72ce 100644
--- a/quickstep/src/com/android/quickstep/util/ActivityInitListener.java
+++ b/quickstep/src/com/android/quickstep/util/ActivityInitListener.java
@@ -31,6 +31,11 @@
     private final BiPredicate<T, Boolean> mOnInitListener;
     private final ActivityTracker<T> mActivityTracker;
 
+    /**
+     * @param onInitListener a callback made when the activity is initialized. The callback should
+     *                       return true to continue receiving callbacks (ie. for if the activity is
+     *                       recreated).
+     */
     public ActivityInitListener(BiPredicate<T, Boolean> onInitListener,
             ActivityTracker<T> tracker) {
         mOnInitListener = onInitListener;
@@ -42,6 +47,10 @@
         return mOnInitListener.test(activity, alreadyOnHome);
     }
 
+    /**
+     * Registers the activity-created listener. If the activity is already created, then the
+     * callback provided in the constructor will be called synchronously.
+     */
     public void register() {
         mActivityTracker.schedule(this);
     }
diff --git a/src/com/android/launcher3/util/ActivityTracker.java b/src/com/android/launcher3/util/ActivityTracker.java
index 4b931f4..499f655 100644
--- a/src/com/android/launcher3/util/ActivityTracker.java
+++ b/src/com/android/launcher3/util/ActivityTracker.java
@@ -50,22 +50,41 @@
         }
     }
 
-    public void schedule(SchedulerCallback<? extends T> callback) {
+    /**
+     * Schedules the callback to be notified when the activity is created.
+     * @return true if the activity is already created, false otherwise
+     */
+    public boolean schedule(SchedulerCallback<? extends T> callback) {
         synchronized (this) {
             mPendingCallback = new WeakReference<>((SchedulerCallback<T>) callback);
         }
-        MAIN_EXECUTOR.execute(this);
+        if (!notifyInitIfPending()) {
+            // If the activity doesn't already exist, then post and wait for the activity to start
+            MAIN_EXECUTOR.execute(this);
+            return false;
+        }
+        return true;
     }
 
     @Override
     public void run() {
-        T activity = mCurrentActivity.get();
-        if (activity != null) {
-            initIfPending(activity, activity.isStarted());
-        }
+        notifyInitIfPending();
     }
 
-    public boolean initIfPending(T activity, boolean alreadyOnHome) {
+    /**
+     * Notifies the pending callback if the activity is now created.
+     * @return true if the activity is now created.
+     */
+    private boolean notifyInitIfPending() {
+        T activity = mCurrentActivity.get();
+        if (activity != null) {
+            notifyInitIfPending(activity, activity.isStarted());
+            return true;
+        }
+        return false;
+    }
+
+    public boolean notifyInitIfPending(T activity, boolean alreadyOnHome) {
         SchedulerCallback<T> pendingCallback = mPendingCallback.get();
         if (pendingCallback != null) {
             if (!pendingCallback.init(activity, alreadyOnHome)) {
@@ -114,7 +133,7 @@
             }
         }
         if (!result && !explicitIntent) {
-            result = initIfPending(activity, alreadyOnHome);
+            result = notifyInitIfPending(activity, alreadyOnHome);
         }
         return result;
     }