diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
index 336cdc9..d84362c 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
@@ -205,15 +205,6 @@
     }
 
     /**
-     * Clean-up logic that occurs when recents is no longer in use/visible.
-     *
-     * @param launcher the launcher activity
-     */
-    public static void resetOverview(Launcher launcher) {
-        launcher.<RecentsView>getOverviewPanel().reset();
-    }
-
-    /**
      * Recents logic that triggers when launcher state changes or launcher activity stops/resumes.
      *
      * @param launcher the launcher activity
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java
index 8436fe5..f1d6450 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java
@@ -21,7 +21,6 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.allapps.AllAppsTransitionController;
 import com.android.launcher3.userevent.nano.LauncherLogProto;
-import com.android.quickstep.util.ClipAnimationHelper;
 import com.android.quickstep.util.LayoutUtils;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
@@ -45,8 +44,6 @@
 
     @Override
     public void onStateEnabled(Launcher launcher) {
-        RecentsView rv = launcher.getOverviewPanel();
-        rv.setOverviewStateEnabled(true);
         AbstractFloatingView.closeAllOpenViews(launcher, false);
     }
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/OverviewState.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/OverviewState.java
index cda03dc..9a99c15 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/OverviewState.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/OverviewState.java
@@ -86,18 +86,10 @@
 
     @Override
     public void onStateEnabled(Launcher launcher) {
-        RecentsView rv = launcher.getOverviewPanel();
-        rv.setOverviewStateEnabled(true);
         AbstractFloatingView.closeAllOpenViews(launcher);
     }
 
     @Override
-    public void onStateDisabled(Launcher launcher) {
-        RecentsView rv = launcher.getOverviewPanel();
-        rv.setOverviewStateEnabled(false);
-    }
-
-    @Override
     public void onStateTransitionEnd(Launcher launcher) {
         launcher.getRotationHelper().setCurrentStateRequest(REQUEST_ROTATE);
         DiscoveryBounce.showForOverviewIfNeeded(launcher);
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java
index 84af109..5104fb8 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java
@@ -36,8 +36,6 @@
  */
 public class TaskOverlayFactory implements ResourceBasedOverride {
 
-    public static final String AIAI_PACKAGE = "com.google.android.as";
-
     /** Note that these will be shown in order from top to bottom, if available for the task. */
     private static final TaskSystemShortcut[] MENU_OPTIONS = new TaskSystemShortcut[]{
             new TaskSystemShortcut.AppInfo(),
@@ -51,29 +49,33 @@
             new MainThreadInitializedObject<>(c -> Overrides.getObject(TaskOverlayFactory.class,
                     c, R.string.task_overlay_factory_class));
 
+    public List<TaskSystemShortcut> getEnabledShortcuts(TaskView taskView) {
+        final ArrayList<TaskSystemShortcut> shortcuts = new ArrayList<>();
+        final BaseDraggingActivity activity = BaseActivity.fromContext(taskView.getContext());
+        for (TaskSystemShortcut menuOption : MENU_OPTIONS) {
+            View.OnClickListener onClickListener =
+                    menuOption.getOnClickListener(activity, taskView);
+            if (onClickListener != null) {
+                shortcuts.add(menuOption);
+            }
+        }
+        return shortcuts;
+    }
+
     public TaskOverlay createOverlay(View thumbnailView) {
         return new TaskOverlay();
     }
 
     public static class TaskOverlay {
 
-        public void setTaskInfo(Task task, ThumbnailData thumbnail, Matrix matrix) {
-        }
+        /**
+         * Called when the current task is interactive for the user
+         */
+        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix) { }
 
-        public void reset() {
-        }
-
-        public List<TaskSystemShortcut> getEnabledShortcuts(TaskView taskView) {
-            final ArrayList<TaskSystemShortcut> shortcuts = new ArrayList<>();
-            final BaseDraggingActivity activity = BaseActivity.fromContext(taskView.getContext());
-            for (TaskSystemShortcut menuOption : MENU_OPTIONS) {
-                View.OnClickListener onClickListener =
-                        menuOption.getOnClickListener(activity, taskView);
-                if (onClickListener != null) {
-                    shortcuts.add(menuOption);
-                }
-            }
-            return shortcuts;
-        }
+        /**
+         * Called when the overlay is no longer used.
+         */
+        public void reset() { }
     }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/fallback/FallbackRecentsView.java
index e254e24..a5aa1bf 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -35,6 +35,7 @@
     public FallbackRecentsView(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
         setOverviewStateEnabled(true);
+        setOverlayEnabled(true);
     }
 
     @Override
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
index deedd21..4162845 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -20,6 +20,7 @@
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.LauncherState.OVERVIEW;
 import static com.android.launcher3.LauncherState.RECENTS_CLEAR_ALL_BUTTON;
+import static com.android.launcher3.LauncherState.SPRING_LOADED;
 import static com.android.launcher3.QuickstepAppTransitionManagerImpl.ALL_APPS_PROGRESS_OFF_SCREEN;
 import static com.android.launcher3.allapps.AllAppsTransitionController.ALL_APPS_PROGRESS;
 import static com.android.launcher3.config.FeatureFlags.ENABLE_QUICKSTEP_LIVE_TILE;
@@ -30,7 +31,6 @@
 import android.content.Context;
 import android.graphics.Canvas;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.os.Build;
 import android.util.AttributeSet;
 import android.view.View;
@@ -40,6 +40,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
+import com.android.launcher3.LauncherStateManager.StateListener;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.appprediction.PredictionUiStateManager;
@@ -57,7 +58,7 @@
  * {@link RecentsView} used in Launcher activity
  */
 @TargetApi(Build.VERSION_CODES.O)
-public class LauncherRecentsView extends RecentsView<Launcher> {
+public class LauncherRecentsView extends RecentsView<Launcher> implements StateListener {
 
     private final TransformParams mTransformParams = new TransformParams();
     private final int mChipOverhang;
@@ -75,6 +76,7 @@
         super(context, attrs, defStyleAttr);
         setContentAlpha(0);
         mChipOverhang = (int) context.getResources().getDimension(R.dimen.chip_hint_overhang);
+        mActivity.getStateManager().addStateListener(this);
     }
 
     @Override
@@ -287,6 +289,20 @@
     }
 
     @Override
+    public void onStateTransitionStart(LauncherState toState) {
+        setOverviewStateEnabled(toState.overviewUi);
+    }
+
+    @Override
+    public void onStateTransitionComplete(LauncherState finalState) {
+        if (finalState == NORMAL || finalState == SPRING_LOADED) {
+            // Clean-up logic that occurs when recents is no longer in use/visible.
+            reset();
+        }
+        setOverlayEnabled(finalState == OVERVIEW);
+    }
+
+    @Override
     public void setOverviewStateEnabled(boolean enabled) {
         super.setOverviewStateEnabled(enabled);
         if (enabled) {
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 1e1007e..38bd461 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
@@ -72,6 +72,9 @@
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.ListView;
 
+import androidx.annotation.Nullable;
+import androidx.dynamicanimation.animation.SpringForce;
+
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Insettable;
@@ -111,9 +114,6 @@
 import java.util.ArrayList;
 import java.util.function.Consumer;
 
-import androidx.annotation.Nullable;
-import androidx.dynamicanimation.animation.SpringForce;
-
 /**
  * A list of recent tasks.
  */
@@ -168,8 +168,6 @@
     // The threshold at which we update the SystemUI flags when animating from the task into the app
     public static final float UPDATE_SYSUI_FLAGS_THRESHOLD = 0.85f;
 
-    private static final float[] sTempFloatArray = new float[3];
-
     protected final T mActivity;
     private final float mFastFlingVelocity;
     private final RecentsModel mModel;
@@ -189,6 +187,7 @@
 
     private boolean mDwbToastShown;
     private boolean mDisallowScrollToClearAll;
+    private boolean mOverlayEnabled;
 
     /**
      * TODO: Call reloadIdNeeded in onTaskStackChanged.
@@ -428,11 +427,7 @@
         // Clear the task data for the removed child if it was visible
         if (child != mClearAllButton) {
             TaskView taskView = (TaskView) child;
-            Task task = taskView.getTask();
-            if (mHasVisibleTaskData.get(task.key.id)) {
-                mHasVisibleTaskData.delete(task.key.id);
-                taskView.onTaskListVisibilityChanged(false /* visible */);
-            }
+            mHasVisibleTaskData.delete(taskView.getTask().key.id);
             mTaskViewPool.recycle(taskView);
         }
     }
@@ -455,7 +450,6 @@
     public void setOverviewStateEnabled(boolean enabled) {
         mOverviewStateEnabled = enabled;
         updateTaskStackListenerState();
-        if (!enabled) mDwbToastShown = false;
     }
 
     public void onDigitalWellbeingToastShown() {
@@ -582,6 +576,7 @@
         }
         resetTaskVisuals();
         onTaskStackUpdated();
+        updateEnabledOverlays();
     }
 
     public int getTaskViewCount() {
@@ -780,6 +775,7 @@
 
         unloadVisibleTaskData();
         setCurrentPage(0);
+        mDwbToastShown = false;
     }
 
     public @Nullable TaskView getRunningTaskView() {
@@ -1545,6 +1541,7 @@
     protected void notifyPageSwitchListener(int prevPage) {
         super.notifyPageSwitchListener(prevPage);
         loadVisibleTaskData();
+        updateEnabledOverlays();
     }
 
     @Override
@@ -1685,4 +1682,19 @@
     public ClipAnimationHelper getTempClipAnimationHelper() {
         return mTempClipAnimationHelper;
     }
+
+    private void updateEnabledOverlays() {
+        int overlayEnabledPage = mOverlayEnabled ? getNextPage() : -1;
+        int taskCount = getTaskViewCount();
+        for (int i = 0; i < taskCount; i++) {
+            ((TaskView) getChildAt(i)).setOverlayEnabled(i == overlayEnabledPage);
+        }
+    }
+
+    public void setOverlayEnabled(boolean overlayEnabled) {
+        if (mOverlayEnabled != overlayEnabled) {
+            mOverlayEnabled = overlayEnabled;
+            updateEnabledOverlays();
+        }
+    }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java
index c47e943..c1f6b82 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java
@@ -42,6 +42,7 @@
 import com.android.launcher3.anim.RoundedRectRevealOutlineProvider;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.BaseDragLayer;
+import com.android.quickstep.TaskOverlayFactory;
 import com.android.quickstep.TaskSystemShortcut;
 import com.android.quickstep.TaskUtils;
 import com.android.quickstep.views.IconView.OnScaleUpdateListener;
@@ -196,7 +197,7 @@
 
         final BaseDraggingActivity activity = BaseDraggingActivity.fromContext(getContext());
         final List<TaskSystemShortcut> shortcuts =
-                taskView.getTaskOverlay().getEnabledShortcuts(taskView);
+                TaskOverlayFactory.INSTANCE.get(getContext()).getEnabledShortcuts(taskView);
         final int count = shortcuts.size();
         for (int i = 0; i < count; ++i) {
             final TaskSystemShortcut menuOption = shortcuts.get(i);
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
index df5831b..8b8240d 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -98,6 +98,9 @@
     private float mDimAlphaMultiplier = 1f;
     private float mSaturation = 1f;
 
+    private boolean mOverlayEnabled;
+    private boolean mRotated;
+
     public TaskThumbnailView(Context context) {
         this(context, null);
     }
@@ -192,10 +195,6 @@
         return 0;
     }
 
-    public TaskOverlay getTaskOverlay() {
-        return mOverlay;
-    }
-
     @Override
     protected void onDraw(Canvas canvas) {
         RectF currentDrawnInsets = mFullscreenParams.mCurrentDrawnInsets;
@@ -257,6 +256,22 @@
         return (TaskView) getParent();
     }
 
+    public void setOverlayEnabled(boolean overlayEnabled) {
+        if (mOverlayEnabled != overlayEnabled) {
+            mOverlayEnabled = overlayEnabled;
+            updateOverlay();
+        }
+    }
+
+    private void updateOverlay() {
+        // The overlay doesn't really work when the screenshot is rotated, so don't add it.
+        if (mOverlayEnabled && !mRotated && mBitmapShader != null && mThumbnailData != null) {
+            mOverlay.initOverlay(mTask, mThumbnailData, mMatrix);
+        } else {
+            mOverlay.reset();
+        }
+    }
+
     private void updateThumbnailPaintFilter() {
         int mul = (int) ((1 - mDimAlpha * mDimAlphaMultiplier) * 255);
         ColorFilter filter = getColorFilter(mul, mIsDarkTextTheme, mSaturation);
@@ -351,12 +366,8 @@
             mPaint.setShader(mBitmapShader);
         }
 
-        if (isRotated) {
-            // The overlay doesn't really work when the screenshot is rotated, so don't add it.
-            mOverlay.reset();
-        } else {
-            mOverlay.setTaskInfo(mTask, mThumbnailData, mMatrix);
-        }
+        mRotated = isRotated;
+        updateOverlay();
         invalidate();
     }
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
index f8d454f..bf3e91f 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
@@ -39,7 +39,6 @@
 import android.util.AttributeSet;
 import android.util.FloatProperty;
 import android.util.Log;
-import android.util.Property;
 import android.view.View;
 import android.view.ViewOutlineProvider;
 import android.view.accessibility.AccessibilityNodeInfo;
@@ -253,10 +252,6 @@
         return mIconView;
     }
 
-    public TaskOverlayFactory.TaskOverlay getTaskOverlay() {
-        return mSnapshotView.getTaskOverlay();
-    }
-
     public AnimatorPlaybackController createLaunchAnimationForRunningTask() {
         final PendingAnimation pendingAnimation =
                 getRecentsView().createTaskLauncherAnimation(this, RECENTS_LAUNCH_DURATION);
@@ -488,6 +483,8 @@
         // Clear any references to the thumbnail (it will be re-read either from the cache or the
         // system on next bind)
         mSnapshotView.setThumbnail(mTask, null);
+        setOverlayEnabled(false);
+        onTaskListVisibilityChanged(false);
         if (mTask != null) {
             mTask.thumbnail = null;
         }
@@ -608,7 +605,7 @@
 
         final Context context = getContext();
         final List<TaskSystemShortcut> shortcuts =
-                mSnapshotView.getTaskOverlay().getEnabledShortcuts(this);
+                TaskOverlayFactory.INSTANCE.get(getContext()).getEnabledShortcuts(this);
         final int count = shortcuts.size();
         for (int i = 0; i < count; ++i) {
             final TaskSystemShortcut menuOption = shortcuts.get(i);
@@ -647,7 +644,7 @@
         }
 
         final List<TaskSystemShortcut> shortcuts =
-                mSnapshotView.getTaskOverlay().getEnabledShortcuts(this);
+                TaskOverlayFactory.INSTANCE.get(getContext()).getEnabledShortcuts(this);
         final int count = shortcuts.size();
         for (int i = 0; i < count; ++i) {
             final TaskSystemShortcut menuOption = shortcuts.get(i);
@@ -733,6 +730,10 @@
         return mShowScreenshot;
     }
 
+    public void setOverlayEnabled(boolean overlayEnabled) {
+        mSnapshotView.setOverlayEnabled(overlayEnabled);
+    }
+
     /**
      * We update and subsequently draw these in {@link #setFullscreenProgress(float)}.
      */
diff --git a/src/com/android/launcher3/LauncherState.java b/src/com/android/launcher3/LauncherState.java
index 3a92dfb..c63f976 100644
--- a/src/com/android/launcher3/LauncherState.java
+++ b/src/com/android/launcher3/LauncherState.java
@@ -258,9 +258,6 @@
      * Called when the start transition ends and the user settles on this particular state.
      */
     public void onStateTransitionEnd(Launcher launcher) {
-        if (this == NORMAL || this == SPRING_LOADED) {
-            UiFactory.resetOverview(launcher);
-        }
         if (this == NORMAL) {
             // Clear any rotation locks when going to normal state
             launcher.getRotationHelper().setCurrentStateRequest(REQUEST_NONE);