diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index b7f3f65..c243a24 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -16,18 +16,22 @@
 
 package com.android.quickstep;
 
+import static com.android.launcher3.Flags.enableRefactorTaskThumbnail;
 import static com.android.quickstep.views.OverviewActionsView.DISABLED_NO_THUMBNAIL;
 import static com.android.quickstep.views.OverviewActionsView.DISABLED_ROTATED;
 
 import android.annotation.SuppressLint;
 import android.content.Context;
+import android.graphics.Bitmap;
 import android.graphics.Insets;
 import android.graphics.Matrix;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.os.Build;
 import android.view.View;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.RequiresApi;
 
 import com.android.launcher3.BaseActivity;
@@ -38,6 +42,7 @@
 import com.android.launcher3.util.ResourceBasedOverride;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.Snackbar;
+import com.android.quickstep.task.util.TaskOverlayHelper;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.quickstep.views.DesktopTaskView;
 import com.android.quickstep.views.GroupedTaskView;
@@ -128,12 +133,43 @@
 
         private T mActionsView;
         protected ImageActionsApi mImageApi;
+        protected TaskOverlayHelper mHelper;
 
         protected TaskOverlay(TaskContainer taskContainer) {
             mApplicationContext = taskContainer.getTaskView().getContext().getApplicationContext();
             mTaskContainer = taskContainer;
-            mImageApi = new ImageActionsApi(
-                    mApplicationContext, mTaskContainer::getThumbnail);
+            if (enableRefactorTaskThumbnail()) {
+                mHelper = new TaskOverlayHelper(mTaskContainer.getTask(), this);
+            }
+            mImageApi = new ImageActionsApi(mApplicationContext, this::getThumbnail);
+        }
+
+        /**
+         * Initialize the overlay when a Task is bound to the TaskView.
+         */
+        public void init() {
+            if (enableRefactorTaskThumbnail()) {
+                mHelper.init();
+            }
+        }
+
+        /**
+         * Destroy the overlay when the TaskView is recycled.
+         */
+        public void destroy() {
+            if (enableRefactorTaskThumbnail()) {
+                mHelper.destroy();
+            }
+        }
+
+        protected @Nullable Bitmap getThumbnail() {
+            return enableRefactorTaskThumbnail() ? mHelper.getEnabledState().getThumbnail()
+                    : mTaskContainer.getThumbnailViewDeprecated().getThumbnail();
+        }
+
+        protected boolean isRealSnapshot() {
+            return enableRefactorTaskThumbnail() ? mHelper.getEnabledState().isRealSnapshot()
+                    : mTaskContainer.getThumbnailViewDeprecated().isRealSnapshot();
         }
 
         protected T getActionsView() {
@@ -151,15 +187,25 @@
 
         /**
          * Called when the current task is interactive for the user
+         *
+         * @deprecated TODO(b/350931107): Remove this interface once TaskOverlayFactoryGo is updated
          */
+        @Deprecated
         public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
                 boolean rotated) {
+            initOverlay(task, thumbnail.getThumbnail(), matrix, rotated);
+        }
+
+        /**
+         * Called when the current task is interactive for the user
+         */
+        public void initOverlay(Task task, @Nullable Bitmap thumbnail, Matrix matrix,
+                boolean rotated) {
             getActionsView().updateDisabledFlags(DISABLED_NO_THUMBNAIL, thumbnail == null);
 
             if (thumbnail != null) {
                 getActionsView().updateDisabledFlags(DISABLED_ROTATED, rotated);
-                boolean isAllowedByPolicy = mTaskContainer.isRealSnapshot();
-                getActionsView().setCallbacks(new OverlayUICallbacksImpl(isAllowedByPolicy, task));
+                getActionsView().setCallbacks(new OverlayUICallbacksImpl(isRealSnapshot(), task));
             }
         }
 
@@ -183,8 +229,8 @@
          */
         @SuppressLint("NewApi")
         protected void saveScreenshot(Task task) {
-            if (mTaskContainer.isRealSnapshot()) {
-                mImageApi.saveScreenshot(mTaskContainer.getThumbnail(),
+            if (isRealSnapshot()) {
+                mImageApi.saveScreenshot(getThumbnail(),
                         getTaskSnapshotBounds(), getTaskSnapshotInsets(), task.key);
             } else {
                 showBlockedByPolicyMessage();
@@ -259,7 +305,36 @@
          */
         @RequiresApi(api = Build.VERSION_CODES.Q)
         public Insets getTaskSnapshotInsets() {
-            return mTaskContainer.getScaledInsets();
+            Bitmap thumbnail = getThumbnail();
+            if (thumbnail == null) {
+                return Insets.NONE;
+            }
+
+            RectF bitmapRect = new RectF(
+                    0,
+                    0,
+                    thumbnail.getWidth(),
+                    thumbnail.getHeight());
+            View snapshotView = mTaskContainer.getSnapshotView();
+            RectF viewRect = new RectF(0, 0, snapshotView.getMeasuredWidth(),
+                    snapshotView.getMeasuredHeight());
+
+            // The position helper matrix tells us how to transform the bitmap to fit the view, the
+            // inverse tells us where the view would be in the bitmaps coordinates. The insets are
+            // the difference between the bitmap bounds and the projected view bounds.
+            Matrix boundsToBitmapSpace = new Matrix();
+            Matrix thumbnailMatrix = enableRefactorTaskThumbnail()
+                    ? mHelper.getEnabledState().getThumbnailMatrix()
+                    : mTaskContainer.getThumbnailViewDeprecated().getThumbnailMatrix();
+            thumbnailMatrix.invert(boundsToBitmapSpace);
+            RectF boundsInBitmapSpace = new RectF();
+            boundsToBitmapSpace.mapRect(boundsInBitmapSpace, viewRect);
+
+            RecentsViewContainer container = RecentsViewContainer.containerFromContext(
+                    getTaskView().getContext());
+            int bottomInset = container.getDeviceProfile().isTablet
+                    ? Math.round(bitmapRect.bottom - boundsInBitmapSpace.bottom) : 0;
+            return Insets.of(0, 0, 0, bottomInset);
         }
 
         /**
diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
index 28212cf..fdb62df 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
@@ -24,4 +24,11 @@
 
     // This is typically a View concern but it is used to invalidate rendering in other Views
     val scale = MutableStateFlow(1f)
+
+    // Whether the current RecentsView state supports task overlays.
+    // TODO(b/331753115): Derive from RecentsView state flow once migrated to MVVM.
+    val overlayEnabled = MutableStateFlow(false)
+
+    // The settled set of visible taskIds that is updated after RecentsView scroll settles.
+    val settledFullyVisibleTaskIds = MutableStateFlow(emptySet<Int>())
 }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
new file mode 100644
index 0000000..feee11f
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 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.quickstep.task.thumbnail
+
+import android.graphics.Bitmap
+import android.graphics.Matrix
+
+/** Ui state for [com.android.quickstep.TaskOverlayFactory.TaskOverlay] */
+sealed class TaskOverlayUiState {
+    data object Disabled : TaskOverlayUiState()
+
+    data class Enabled(
+        val isRealSnapshot: Boolean,
+        val thumbnail: Bitmap?,
+        val thumbnailMatrix: Matrix
+    ) : TaskOverlayUiState()
+}
diff --git a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
new file mode 100644
index 0000000..de39584
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2024 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.quickstep.task.util
+
+import android.util.Log
+import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
+import com.android.quickstep.task.viewmodel.TaskOverlayViewModel
+import com.android.quickstep.views.RecentsView
+import com.android.quickstep.views.RecentsViewContainer
+import com.android.systemui.shared.recents.model.Task
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.MainScope
+import kotlinx.coroutines.launch
+
+/**
+ * Helper for [TaskOverlayFactory.TaskOverlay] to interact with [TaskOverlayViewModel], this helper
+ * should merge with [TaskOverlayFactory.TaskOverlay] when it's migrated to MVVM.
+ */
+class TaskOverlayHelper(val task: Task, val overlay: TaskOverlayFactory.TaskOverlay<*>) {
+    private lateinit var job: Job
+    private var uiState: TaskOverlayUiState = Disabled
+
+    // TODO(b/335649589): Ideally create and obtain this from DI. This ViewModel should be scoped
+    //  to [TaskView], and also shared between [TaskView] and [TaskThumbnailView]
+    //  This is using a lazy for now because the dependencies cannot be obtained without DI.
+    private val taskOverlayViewModel by lazy {
+        val recentsView =
+            RecentsViewContainer.containerFromContext<RecentsViewContainer>(
+                    overlay.taskView.context
+                )
+                .getOverviewPanel<RecentsView<*, *>>()
+        TaskOverlayViewModel(task, recentsView.mRecentsViewData, recentsView.mTasksRepository)
+    }
+
+    // TODO(b/331753115): TaskOverlay should listen for state changes and react.
+    val enabledState: Enabled
+        get() = uiState as Enabled
+
+    fun init() {
+        // TODO(b/335396935): This should be changed to TaskView's scope.
+        job =
+            MainScope().launch {
+                taskOverlayViewModel.overlayState.collect {
+                    uiState = it
+                    if (it is Enabled) {
+                        Log.d(
+                            TAG,
+                            "initOverlay - taskId: ${task.key.id}, thumbnail: ${it.thumbnail}"
+                        )
+                        overlay.initOverlay(
+                            task,
+                            it.thumbnail,
+                            it.thumbnailMatrix,
+                            /* rotated= */ false
+                        )
+                    } else {
+                        Log.d(TAG, "reset - taskId: ${task.key.id}")
+                        overlay.reset()
+                    }
+                }
+            }
+    }
+
+    fun destroy() {
+        job.cancel()
+        uiState = Disabled
+        overlay.reset()
+    }
+
+    companion object {
+        private const val TAG = "TaskOverlayHelper"
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
new file mode 100644
index 0000000..4682323
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2024 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.quickstep.task.viewmodel
+
+import android.graphics.Matrix
+import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
+import com.android.systemui.shared.recents.model.Task
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.distinctUntilChangedBy
+import kotlinx.coroutines.flow.map
+
+/** View model for TaskOverlay */
+class TaskOverlayViewModel(
+    task: Task,
+    recentsViewData: RecentsViewData,
+    tasksRepository: RecentTasksRepository,
+) {
+    val overlayState =
+        combine(
+                recentsViewData.overlayEnabled,
+                recentsViewData.settledFullyVisibleTaskIds.map { it.contains(task.key.id) },
+                tasksRepository
+                    .getTaskDataById(task.key.id)
+                    .map { it?.thumbnail }
+                    .distinctUntilChangedBy { it?.snapshotId }
+            ) { isOverlayEnabled, isFullyVisible, thumbnailData ->
+                if (isOverlayEnabled && isFullyVisible) {
+                    Enabled(
+                        isRealSnapshot = (thumbnailData?.isRealSnapshot ?: false) && !task.isLocked,
+                        thumbnailData?.thumbnail,
+                        // TODO(b/343101424): Use PreviewPositionHelper, listen from a common source
+                        // with
+                        //  TaskThumbnailView.
+                        Matrix.IDENTITY_MATRIX
+                    )
+                } else {
+                    Disabled
+                }
+            }
+            .distinctUntilChanged()
+}
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 8553635..fbc7283 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -229,10 +229,12 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -2390,12 +2392,11 @@
             if (containers.isEmpty()) {
                 continue;
             }
-            int index = indexOfChild(taskView);
             boolean visible;
             if (showAsGrid()) {
                 visible = isTaskViewWithinBounds(taskView, visibleStart, visibleEnd);
             } else {
-                visible = lower <= index && index <= upper;
+                visible = lower <= i && i <= upper;
             }
             if (visible) {
                 // Default update all non-null tasks, then remove running ones
@@ -2426,7 +2427,7 @@
                         }
                         taskView.onTaskListVisibilityChanged(true /* visible */, changes);
                     }
-                    mHasVisibleTaskData.put(task.key.id, visible);
+                    mHasVisibleTaskData.put(task.key.id, true);
                 }
             } else {
                 for (TaskContainer container : containers) {
@@ -5346,6 +5347,43 @@
         updateCurrentTaskActionsVisibility();
         loadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
         updateEnabledOverlays();
+
+        if (enableRefactorTaskThumbnail()) {
+            int screenStart = 0;
+            int screenEnd = 0;
+            int centerPageIndex = 0;
+            if (showAsGrid()) {
+                screenStart = getPagedOrientationHandler().getPrimaryScroll(this);
+                int pageOrientedSize = getPagedOrientationHandler().getMeasuredSize(this);
+                screenEnd = screenStart + pageOrientedSize;
+            } else {
+                centerPageIndex = getPageNearestToCenterOfScreen();
+            }
+
+            Set<Integer> fullyVisibleTaskIds = new HashSet<>();
+
+            // Update the task data for the in/visible children
+            for (int i = 0; i < getTaskViewCount(); i++) {
+                TaskView taskView = requireTaskViewAt(i);
+                List<TaskContainer> containers = taskView.getTaskContainers();
+                if (containers.isEmpty()) {
+                    continue;
+                }
+                boolean isFullyVisible;
+                if (showAsGrid()) {
+                    isFullyVisible = isTaskViewFullyWithinBounds(taskView, screenStart,
+                            screenEnd);
+                } else {
+                    isFullyVisible = i == centerPageIndex;
+                }
+                if (isFullyVisible) {
+                    List<Integer> taskIds = containers.stream().map(
+                            taskContainer -> taskContainer.getTask().key.id).toList();
+                    fullyVisibleTaskIds.addAll(taskIds);
+                }
+            }
+            mRecentsViewData.getSettledFullyVisibleTaskIds().setValue(fullyVisibleTaskIds);
+        }
     }
 
     @Override
@@ -5896,6 +5934,7 @@
         if (mOverlayEnabled != overlayEnabled) {
             mOverlayEnabled = overlayEnabled;
             updateEnabledOverlays();
+            mRecentsViewData.getOverlayEnabled().setValue(overlayEnabled);
         }
     }
 
@@ -6034,7 +6073,7 @@
      * tasks to be dimmed while other elements in the recents view are left alone.
      */
     public void showForegroundScrim(boolean show) {
-        // TODO(b/335606129) Add scrim response into new TTV - this is called from overlay
+        // TODO(b/349601769) Add scrim response into new TTV - this is called from overlay
         if (!show && mColorTint == 0) {
             if (mTintingAnimator != null) {
                 mTintingAnimator.cancel();
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index cfdee6c..5c6dd0c 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -18,9 +18,9 @@
 
 import android.content.Intent
 import android.graphics.Bitmap
-import android.graphics.Insets
 import android.view.View
-import com.android.launcher3.Flags
+import com.android.launcher3.Flags.enableRefactorTaskThumbnail
+import com.android.launcher3.Flags.privateSpaceRestrictAccessibilityDrag
 import com.android.launcher3.LauncherSettings
 import com.android.launcher3.model.data.ItemInfoWithIcon
 import com.android.launcher3.model.data.WorkspaceItemInfo
@@ -64,14 +64,6 @@
     val thumbnail: Bitmap?
         get() = thumbnailViewDeprecated.thumbnail
 
-    // TODO(b/349120849): Extract ThumbnailData from TaskContainerData/TaskThumbnailViewModel
-    val isRealSnapshot: Boolean
-        get() = thumbnailViewDeprecated.isRealSnapshot()
-
-    // TODO(b/349120849): Extract ThumbnailData from TaskContainerData/TaskThumbnailViewModel
-    val scaledInsets: Insets
-        get() = thumbnailViewDeprecated.scaledInsets
-
     /** Builds proto for logging */
     val itemInfo: WorkspaceItemInfo
         get() =
@@ -83,7 +75,7 @@
                 intent = Intent().setComponent(componentKey.componentName)
                 title = task.title
                 taskView.recentsView?.let { screenId = it.indexOfChild(taskView) }
-                if (Flags.privateSpaceRestrictAccessibilityDrag()) {
+                if (privateSpaceRestrictAccessibilityDrag()) {
                     if (
                         UserCache.getInstance(taskView.context)
                             .getUserInfo(componentKey.user)
@@ -98,12 +90,14 @@
     fun destroy() {
         digitalWellBeingToast?.destroy()
         thumbnailView?.let { taskView.removeView(it) }
+        overlay.destroy()
     }
 
     fun bind() {
-        if (Flags.enableRefactorTaskThumbnail() && thumbnailView != null) {
+        if (enableRefactorTaskThumbnail() && thumbnailView != null) {
             thumbnailViewDeprecated.setTaskOverlay(overlay)
             bindThumbnailView()
+            overlay.init()
         } else {
             thumbnailViewDeprecated.bind(task, overlay)
         }
@@ -116,4 +110,10 @@
         //  this should be decided inside TaskThumbnailViewModel.
         thumbnailView?.viewModel?.bind(TaskThumbnail(task.key.id, taskView.isRunningTask))
     }
+
+    fun setOverlayEnabled(enabled: Boolean) {
+        if (!enableRefactorTaskThumbnail()) {
+            thumbnailViewDeprecated.setOverlayEnabled(enabled)
+        }
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
index 2afb6a6..5b7e6c7 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
@@ -28,16 +28,13 @@
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.ColorFilter;
-import android.graphics.Insets;
 import android.graphics.Matrix;
 import android.graphics.Paint;
 import android.graphics.PorterDuff;
 import android.graphics.PorterDuffXfermode;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.graphics.Shader;
 import android.graphics.drawable.Drawable;
-import android.os.Build;
 import android.util.AttributeSet;
 import android.util.FloatProperty;
 import android.util.Property;
@@ -45,7 +42,6 @@
 import android.widget.ImageView;
 
 import androidx.annotation.Nullable;
-import androidx.annotation.RequiresApi;
 import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.DeviceProfile;
@@ -266,40 +262,6 @@
         return mDimAlpha;
     }
 
-    /**
-     * Get the scaled insets that are being used to draw the task view. This is a subsection of
-     * the full snapshot.
-     *
-     * @return the insets in snapshot bitmap coordinates.
-     */
-    @RequiresApi(api = Build.VERSION_CODES.Q)
-    public Insets getScaledInsets() {
-        if (mThumbnailData == null) {
-            return Insets.NONE;
-        }
-
-        RectF bitmapRect = new RectF(
-                0,
-                0,
-                mThumbnailData.getThumbnail().getWidth(),
-                mThumbnailData.getThumbnail().getHeight());
-        RectF viewRect = new RectF(0, 0, getMeasuredWidth(), getMeasuredHeight());
-
-        // The position helper matrix tells us how to transform the bitmap to fit the view, the
-        // inverse tells us where the view would be in the bitmaps coordinates. The insets are the
-        // difference between the bitmap bounds and the projected view bounds.
-        Matrix boundsToBitmapSpace = new Matrix();
-        mPreviewPositionHelper.getMatrix().invert(boundsToBitmapSpace);
-        RectF boundsInBitmapSpace = new RectF();
-        boundsToBitmapSpace.mapRect(boundsInBitmapSpace, viewRect);
-
-        DeviceProfile dp = mContainer.getDeviceProfile();
-        int bottomInset = dp.isTablet
-                ? Math.round(bitmapRect.bottom - boundsInBitmapSpace.bottom) : 0;
-        return Insets.of(0, 0, 0, bottomInset);
-    }
-
-
     @SystemUiControllerFlags
     public int getSysUiStatusNavFlags() {
         if (mThumbnailData != null) {
@@ -487,7 +449,9 @@
      */
     private void refreshOverlay() {
         if (mOverlayEnabled) {
-            mOverlay.initOverlay(mTask, mThumbnailData, mPreviewPositionHelper.getMatrix(),
+            mOverlay.initOverlay(mTask,
+                    mThumbnailData != null ? mThumbnailData.getThumbnail() : null,
+                    mPreviewPositionHelper.getMatrix(),
                     mPreviewPositionHelper.isOrientationChanged());
         } else {
             mOverlay.reset();
@@ -560,6 +524,10 @@
         return mThumbnailData.isRealSnapshot && !mTask.isLocked;
     }
 
+    public Matrix getThumbnailMatrix() {
+        return mPreviewPositionHelper.getMatrix();
+    }
+
     @Override
     public void onRecycle() {
         // Do nothing
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index b922df4..5951cdc 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -1394,10 +1394,8 @@
     }
 
     open fun setOverlayEnabled(overlayEnabled: Boolean) {
-        // TODO(b/335606129) Investigate the usage of [TaskOverlay] in the new TaskThumbnailView.
-        //  and if it's still necessary we should support that in the new TTV class.
         if (!enableRefactorTaskThumbnail()) {
-            taskContainers.forEach { it.thumbnailViewDeprecated.setOverlayEnabled(overlayEnabled) }
+            taskContainers.forEach { it.setOverlayEnabled(overlayEnabled) }
         }
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
new file mode 100644
index 0000000..40482c4
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
@@ -0,0 +1,160 @@
+/*
+ * Copyright (C) 2024 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.quickstep.task.viewmodel
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Bitmap
+import android.graphics.Color
+import android.graphics.Matrix
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [TaskOverlayViewModel] */
+@RunWith(AndroidJUnit4::class)
+class TaskOverlayViewModelTest {
+    private val task =
+        Task(Task.TaskKey(TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
+            colorBackground = Color.BLACK
+        }
+    private val thumbnailData =
+        ThumbnailData(
+            thumbnail =
+                mock<Bitmap>().apply {
+                    whenever(width).thenReturn(THUMBNAIL_WIDTH)
+                    whenever(height).thenReturn(THUMBNAIL_HEIGHT)
+                }
+        )
+    private val recentsViewData = RecentsViewData()
+    private val tasksRepository = FakeTasksRepository()
+    private val systemUnderTest = TaskOverlayViewModel(task, recentsViewData, tasksRepository)
+
+    @Test
+    fun initialStateIsDisabled() = runTest {
+        assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
+    }
+
+    @Test
+    fun recentsViewOverlayDisabled_Disabled() = runTest {
+        recentsViewData.overlayEnabled.value = false
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+
+        assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
+    }
+
+    @Test
+    fun taskNotFullyVisible_Disabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf()
+
+        assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
+    }
+
+    @Test
+    fun noThumbnail_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        task.isLocked = false
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = false,
+                    thumbnail = null,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    @Test
+    fun withThumbnail_RealSnapshot_NotLocked_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+        thumbnailData.isRealSnapshot = true
+        task.isLocked = false
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = true,
+                    thumbnail = thumbnailData.thumbnail,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    @Test
+    fun withThumbnail_RealSnapshot_Locked_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+        thumbnailData.isRealSnapshot = true
+        task.isLocked = true
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = false,
+                    thumbnail = thumbnailData.thumbnail,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    @Test
+    fun withThumbnail_FakeSnapshot_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+        thumbnailData.isRealSnapshot = false
+        task.isLocked = false
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = false,
+                    thumbnail = thumbnailData.thumbnail,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    companion object {
+        const val TASK_ID = 0
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
+}