diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index deb06c9..df66a5e 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -42,8 +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.thumbnail.TaskOverlayUiState;
-import com.android.quickstep.task.util.TaskOverlayHelper;
+import com.android.quickstep.recents.domain.usecase.ThumbnailPosition;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.quickstep.views.DesktopTaskView;
 import com.android.quickstep.views.GroupedTaskView;
@@ -53,6 +52,7 @@
 import com.android.quickstep.views.TaskContainer;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -135,54 +135,37 @@
 
         private T mActionsView;
         protected ImageActionsApi mImageApi;
-        protected TaskOverlayHelper mHelper;
+        private ThumbnailData mThumbnailData = null;
 
         protected TaskOverlay(TaskContainer taskContainer) {
             mApplicationContext = taskContainer.getTaskView().getContext().getApplicationContext();
             mTaskContainer = taskContainer;
-            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();
-            }
+        public void setThumbnailState(@Nullable ThumbnailData thumbnailData) {
+            mThumbnailData = thumbnailData;
         }
 
         protected @Nullable Bitmap getThumbnail() {
-            return enableRefactorTaskThumbnail() ? mHelper.getEnabledState().getThumbnail()
-                    : mTaskContainer.getThumbnailViewDeprecated().getThumbnail();
+            if (enableRefactorTaskThumbnail()) {
+                return mThumbnailData == null ? null : mThumbnailData.getThumbnail();
+            } else {
+                return mTaskContainer.getThumbnailViewDeprecated().getThumbnail();
+            }
         }
-
         /**
          * Returns whether the snapshot is real. If the device is locked for the user of the task,
          * the snapshot used will be an app-theme generated snapshot instead of a real snapshot.
          */
-        public boolean isRealSnapshot() {
+        protected boolean isRealSnapshot() {
             if (enableRefactorTaskThumbnail()) {
-                if (mHelper.getUiState() instanceof TaskOverlayUiState.Enabled) {
-                    return mHelper.getEnabledState().isRealSnapshot();
-                } else {
-                    return false;
-                }
-            }
+                if (mThumbnailData == null) return false;
 
-            return mTaskContainer.getThumbnailViewDeprecated().isRealSnapshot();
+                return mThumbnailData.isRealSnapshot && !mTaskContainer.getTask().isLocked;
+            } else {
+                return mTaskContainer.getThumbnailViewDeprecated().isRealSnapshot();
+            }
         }
 
         /**
@@ -190,7 +173,8 @@
          */
         public boolean isThumbnailRotationDifferentFromTask() {
             if (enableRefactorTaskThumbnail()) {
-                return mHelper.getThumbnailPositionState().isRotated();
+                ThumbnailPosition thumbnailPosition = mTaskContainer.getThumbnailPosition();
+                return thumbnailPosition != null && thumbnailPosition.isRotated();
             }
 
             return mTaskContainer.getThumbnailViewDeprecated()
@@ -340,9 +324,16 @@
             // 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.getThumbnailMatrix()
-                    : mTaskContainer.getThumbnailViewDeprecated().getThumbnailMatrix();
+            Matrix thumbnailMatrix;
+            if (enableRefactorTaskThumbnail()) {
+                if (mTaskContainer.getThumbnailPosition() != null) {
+                    thumbnailMatrix = mTaskContainer.getThumbnailPosition().getMatrix();
+                } else {
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX;
+                }
+            } else {
+                thumbnailMatrix = mTaskContainer.getThumbnailViewDeprecated().getThumbnailMatrix();
+            }
             thumbnailMatrix.invert(boundsToBitmapSpace);
             RectF boundsInBitmapSpace = new RectF();
             boundsToBitmapSpace.mapRect(boundsInBitmapSpace, viewRect);
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
index 118a931..8a6a805 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
@@ -36,6 +36,7 @@
     val isLiveTile: Boolean,
     val hasHeader: Boolean,
     val sysUiStatusNavFlags: Int,
+    val taskOverlayEnabled: Boolean,
 )
 
 sealed class TaskData {
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
index 3c4a384..09e2071 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
@@ -29,6 +29,7 @@
 import com.android.quickstep.recents.domain.usecase.ThumbnailPosition
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.views.TaskViewType
+import com.android.quickstep.views.TaskViewType.SINGLE
 import com.android.systemui.shared.recents.model.ThumbnailData
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
@@ -74,8 +75,18 @@
             )
         }
 
+    private val overlayEnabled =
+        combine(recentsViewData.overlayEnabled, recentsViewData.settledFullyVisibleTaskIds) {
+                isOverlayEnabled,
+                settledFullyVisibleTaskIds ->
+                taskViewType == SINGLE &&
+                    isOverlayEnabled &&
+                    settledFullyVisibleTaskIds.any { it in taskIds.value }
+            }
+            .distinctUntilChanged()
+
     val state: Flow<TaskTileUiState> =
-        combine(taskData, isLiveTile) { tasks, isLiveTile -> mapToTaskTile(tasks, isLiveTile) }
+        combine(taskData, isLiveTile, overlayEnabled, ::mapToTaskTile)
             .distinctUntilChanged()
             .flowOn(dispatcherProvider.background)
 
@@ -99,13 +110,18 @@
             isRtl = isRtl,
         )
 
-    private fun mapToTaskTile(tasks: List<TaskData>, isLiveTile: Boolean): TaskTileUiState {
+    private fun mapToTaskTile(
+        tasks: List<TaskData>,
+        isLiveTile: Boolean,
+        overlayEnabled: Boolean,
+    ): TaskTileUiState {
         val firstThumbnailData = (tasks.firstOrNull() as? TaskData.Data)?.thumbnailData
         return TaskTileUiState(
             tasks = tasks,
             isLiveTile = isLiveTile,
             hasHeader = taskViewType == TaskViewType.DESKTOP,
             sysUiStatusNavFlags = getSysUiStatusNavFlagsUseCase(firstThumbnailData),
+            taskOverlayEnabled = overlayEnabled,
         )
     }
 
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
deleted file mode 100644
index 5fb5b90..0000000
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * 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
-
-/** Ui state for [com.android.quickstep.TaskOverlayFactory.TaskOverlay] */
-sealed class TaskOverlayUiState {
-    data object Disabled : TaskOverlayUiState()
-
-    data class Enabled(val isRealSnapshot: Boolean, val thumbnail: Bitmap?) : TaskOverlayUiState()
-}
diff --git a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
deleted file mode 100644
index d8aea9d..0000000
--- a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
+++ /dev/null
@@ -1,128 +0,0 @@
-/*
- * 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 android.view.View.OnLayoutChangeListener
-import com.android.launcher3.util.coroutines.DispatcherProvider
-import com.android.quickstep.TaskOverlayFactory
-import com.android.quickstep.recents.di.RecentsDependencies
-import com.android.quickstep.recents.di.get
-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.systemui.shared.recents.model.Task
-import kotlinx.coroutines.CoroutineName
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
-import kotlinx.coroutines.flow.dropWhile
-import kotlinx.coroutines.flow.flowOn
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
-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 val scope = overlay.taskView.context
-    private val recentsCoroutineScope: CoroutineScope = RecentsDependencies.get(scope)
-    private val dispatcherProvider: DispatcherProvider = RecentsDependencies.get(scope)
-    private lateinit var overlayInitializedScope: CoroutineScope
-    var uiState: TaskOverlayUiState = Disabled
-
-    private lateinit var viewModel: TaskOverlayViewModel
-
-    // TODO(b/331753115): TaskOverlay should listen for state changes and react.
-    val enabledState: Enabled
-        get() = uiState as Enabled
-
-    private val snapshotLayoutChangeListener = OnLayoutChangeListener { _, _, _, _, _, _, _, _, _ ->
-        (uiState as? Enabled)?.let { initOverlay(it) }
-    }
-
-    fun getThumbnailMatrix() = getThumbnailPositionState().matrix
-
-    fun getThumbnailPositionState() =
-        viewModel.getThumbnailPositionState(
-            overlay.snapshotView.width,
-            overlay.snapshotView.height,
-            overlay.snapshotView.isLayoutRtl,
-        )
-
-    fun init() {
-        overlayInitializedScope =
-            CoroutineScope(
-                SupervisorJob() + Dispatchers.Main.immediate + CoroutineName("TaskOverlayHelper")
-            )
-        viewModel =
-            TaskOverlayViewModel(
-                task = task,
-                recentsViewData = RecentsDependencies.get(scope),
-                getThumbnailPositionUseCase = RecentsDependencies.get(scope),
-                recentTasksRepository = RecentsDependencies.get(scope),
-                dispatcherProvider = RecentsDependencies.get(scope),
-            )
-        viewModel.overlayState
-            .dropWhile { it == Disabled }
-            .flowOn(dispatcherProvider.background)
-            .onEach {
-                uiState = it
-                if (it is Enabled) {
-                    initOverlay(it)
-                } else {
-                    reset()
-                }
-            }
-            .launchIn(overlayInitializedScope)
-        overlay.snapshotView.addOnLayoutChangeListener(snapshotLayoutChangeListener)
-    }
-
-    private fun initOverlay(enabledState: Enabled) {
-        if (DEBUG) {
-            Log.d(TAG, "initOverlay - taskId: ${task.key.id}, thumbnail: ${enabledState.thumbnail}")
-        }
-        with(getThumbnailPositionState()) {
-            overlay.initOverlay(task, enabledState.thumbnail, matrix, isRotated)
-        }
-    }
-
-    private fun reset() {
-        if (DEBUG) {
-            Log.d(TAG, "reset - taskId: ${task.key.id}")
-        }
-        overlay.reset()
-    }
-
-    fun destroy() {
-        recentsCoroutineScope.launch(dispatcherProvider.background) {
-            overlayInitializedScope.cancel("TaskOverlay being destroyed")
-        }
-        uiState = Disabled
-        overlay.snapshotView.removeOnLayoutChangeListener(snapshotLayoutChangeListener)
-        reset()
-    }
-
-    companion object {
-        private const val TAG = "TaskOverlayHelper"
-        private const val DEBUG = false
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
deleted file mode 100644
index 9bff3ac..0000000
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * 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.launcher3.util.coroutines.DispatcherProvider
-import com.android.quickstep.recents.data.RecentTasksRepository
-import com.android.quickstep.recents.domain.usecase.GetThumbnailPositionUseCase
-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.flowOn
-import kotlinx.coroutines.flow.map
-
-/** View model for TaskOverlay */
-class TaskOverlayViewModel(
-    private val task: Task,
-    recentsViewData: RecentsViewData,
-    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
-    private val recentTasksRepository: RecentTasksRepository,
-    dispatcherProvider: DispatcherProvider,
-) {
-    val overlayState =
-        combine(
-                recentsViewData.overlayEnabled,
-                recentsViewData.settledFullyVisibleTaskIds
-                    .map { it.contains(task.key.id) }
-                    .distinctUntilChanged(),
-                recentTasksRepository.getThumbnailById(task.key.id),
-            ) { isOverlayEnabled, isFullyVisible, thumbnailData ->
-                if (isOverlayEnabled && isFullyVisible) {
-                    Enabled(
-                        isRealSnapshot = (thumbnailData?.isRealSnapshot ?: false) && !task.isLocked,
-                        thumbnailData?.thumbnail,
-                    )
-                } else {
-                    Disabled
-                }
-            }
-            .distinctUntilChanged()
-            .flowOn(dispatcherProvider.background)
-
-    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
-        val thumbnailPositionState =
-            getThumbnailPositionUseCase(
-                thumbnailData = recentTasksRepository.getCurrentThumbnailById(task.key.id),
-                width = width,
-                height = height,
-                isRtl = isRtl,
-            )
-        return ThumbnailPositionState(
-            thumbnailPositionState.matrix,
-            thumbnailPositionState.isRotated,
-        )
-    }
-
-    data class ThumbnailPositionState(val matrix: Matrix, val isRotated: Boolean)
-}
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 6d7ae70..ec6d1c4 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -18,6 +18,7 @@
 
 import android.graphics.Bitmap
 import android.graphics.Matrix
+import android.util.Log
 import android.view.View
 import android.view.View.OnClickListener
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
@@ -26,6 +27,7 @@
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory
 import com.android.quickstep.ViewUtils.addAccessibleChildToList
+import com.android.quickstep.recents.domain.usecase.ThumbnailPosition
 import com.android.quickstep.recents.ui.mapper.TaskUiStateMapper
 import com.android.quickstep.recents.ui.viewmodel.TaskData
 import com.android.quickstep.task.thumbnail.TaskContentView
@@ -54,6 +56,8 @@
     taskOverlayFactory: TaskOverlayFactory,
 ) {
     val overlay: TaskOverlayFactory.TaskOverlay<*> = taskOverlayFactory.createOverlay(this)
+    var thumbnailPosition: ThumbnailPosition? = null
+    private var overlayEnabledStatus = false
 
     init {
         if (enableRefactorTaskThumbnail()) {
@@ -101,14 +105,13 @@
         if (!enableRefactorTaskThumbnail()) {
             thumbnailViewDeprecated.bind(task, overlay, taskView)
         }
-        overlay.init()
     }
 
     fun destroy() {
         digitalWellBeingToast?.destroy()
         taskContentView.scaleX = 1f
         taskContentView.scaleY = 1f
-        overlay.destroy()
+        overlay.reset()
         if (enableRefactorTaskThumbnail()) {
             isThumbnailValid = false
             thumbnailData = null
@@ -124,6 +127,34 @@
         }
     }
 
+    fun setOverlayEnabled(enabled: Boolean, thumbnailPosition: ThumbnailPosition?) {
+        if (enableRefactorTaskThumbnail()) {
+            if (overlayEnabledStatus != enabled || this.thumbnailPosition != thumbnailPosition) {
+                overlayEnabledStatus = enabled
+
+                refreshOverlay(thumbnailPosition)
+            }
+        }
+    }
+
+    fun refreshOverlay(thumbnailPosition: ThumbnailPosition?) {
+        this.thumbnailPosition = thumbnailPosition
+        when {
+            !overlayEnabledStatus -> overlay.reset()
+            thumbnailPosition == null -> {
+                Log.e(TAG, "Thumbnail position was null during overlay refresh", Exception())
+                overlay.reset()
+            }
+            else ->
+                overlay.initOverlay(
+                    task,
+                    thumbnailData?.thumbnail,
+                    thumbnailPosition.matrix,
+                    thumbnailPosition.isRotated,
+                )
+        }
+    }
+
     fun addChildForAccessibility(outChildren: ArrayList<View>) {
         addAccessibleChildToList(iconView.asView(), outChildren)
         addAccessibleChildToList(snapshotView, outChildren)
@@ -144,6 +175,7 @@
             state?.taskId,
         )
         thumbnailData = if (state is TaskData.Data) state.thumbnailData else null
+        overlay.setThumbnailState(thumbnailData)
     }
 
     fun updateTintAmount(tintAmount: Float) {
@@ -184,4 +216,8 @@
     fun updateThumbnailMatrix(matrix: Matrix) {
         thumbnailView.setImageMatrix(matrix)
     }
+
+    companion object {
+        const val TAG = "TaskContainer"
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index c878d56..b7f1d1d 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -82,6 +82,7 @@
 import com.android.quickstep.recents.di.RecentsDependencies
 import com.android.quickstep.recents.di.get
 import com.android.quickstep.recents.di.inject
+import com.android.quickstep.recents.domain.usecase.ThumbnailPosition
 import com.android.quickstep.recents.ui.viewmodel.TaskData
 import com.android.quickstep.recents.ui.viewmodel.TaskTileUiState
 import com.android.quickstep.recents.ui.viewmodel.TaskViewModel
@@ -775,11 +776,13 @@
                     },
             )
             updateThumbnailValidity(container)
-            updateThumbnailMatrix(
-                container = container,
-                width = container.thumbnailView.width,
-                height = container.thumbnailView.height,
-            )
+            val thumbnailPosition =
+                updateThumbnailMatrix(
+                    container = container,
+                    width = container.thumbnailView.width,
+                    height = container.thumbnailView.height,
+                )
+            container.setOverlayEnabled(state.taskOverlayEnabled, thumbnailPosition)
 
             if (enableOverviewIconMenu()) {
                 setIconState(container, containerState)
@@ -808,11 +811,16 @@
      * @param width The desired width of the thumbnail's container.
      * @param height The desired height of the thumbnail's container.
      */
-    private fun updateThumbnailMatrix(container: TaskContainer, width: Int, height: Int) {
+    private fun updateThumbnailMatrix(
+        container: TaskContainer,
+        width: Int,
+        height: Int,
+    ): ThumbnailPosition? {
         val thumbnailPosition =
             viewModel?.getThumbnailPosition(container.thumbnailData, width, height, isLayoutRtl)
-                ?: return
+                ?: return null
         container.updateThumbnailMatrix(thumbnailPosition.matrix)
+        return thumbnailPosition
     }
 
     override fun onDetachedFromWindow() {
@@ -877,7 +885,8 @@
                     thumbnailFullscreenParams.currentCornerRadius
                 container.taskContentView.doOnSizeChange { width, height ->
                     updateThumbnailValidity(container)
-                    updateThumbnailMatrix(container, width, height)
+                    val thumbnailPosition = updateThumbnailMatrix(container, width, height)
+                    container.refreshOverlay(thumbnailPosition)
                 }
             }
         }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
index a97ef0c..18b9fe9 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
@@ -66,16 +66,7 @@
 
     @Before
     fun setUp() {
-        sut =
-            TaskViewModel(
-                taskViewType = TaskViewType.SINGLE,
-                recentsViewData = recentsViewData,
-                getTaskUseCase = getTaskUseCase,
-                getSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase(),
-                isThumbnailValidUseCase = isThumbnailValidUseCase,
-                getThumbnailPositionUseCase = getThumbnailPositionUseCase,
-                dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
-            )
+        sut = createTaskViewModel(TaskViewType.SINGLE)
         whenever(getTaskUseCase.invoke(TASK_MODEL_1.id)).thenReturn(flow { emit(TASK_MODEL_1) })
         whenever(getTaskUseCase.invoke(TASK_MODEL_2.id)).thenReturn(flow { emit(TASK_MODEL_2) })
         whenever(getTaskUseCase.invoke(TASK_MODEL_3.id)).thenReturn(flow { emit(TASK_MODEL_3) })
@@ -93,6 +84,7 @@
                     isLiveTile = false,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -139,6 +131,7 @@
                     isLiveTile = false,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -161,6 +154,7 @@
                     isLiveTile = true,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -183,6 +177,7 @@
                     isLiveTile = false,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -204,6 +199,7 @@
                     isLiveTile = false,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -221,6 +217,7 @@
                     isLiveTile = false,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -235,11 +232,64 @@
                     isLiveTile = false,
                     hasHeader = false,
                     sysUiStatusNavFlags = FLAGS_APPEARANCE_DEFAULT,
+                    taskOverlayEnabled = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
 
     @Test
+    fun taskOverlayEnabled_when_OverlayIsEnabledForVisibleSingleTask() =
+        testScope.runTest {
+            sut.bind(TASK_MODEL_1.id)
+            recentsViewData.overlayEnabled.value = true
+            recentsViewData.settledFullyVisibleTaskIds.value = setOf(1)
+
+            assertThat(sut.state.first().taskOverlayEnabled).isTrue()
+        }
+
+    @Test
+    fun taskOverlayDisabled_when_usingGroupedTask() =
+        testScope.runTest {
+            sut = createTaskViewModel(TaskViewType.GROUPED)
+            sut.bind(TASK_MODEL_1.id)
+            recentsViewData.overlayEnabled.value = true
+            recentsViewData.settledFullyVisibleTaskIds.value = setOf(1)
+
+            assertThat(sut.state.first().taskOverlayEnabled).isFalse()
+        }
+
+    @Test
+    fun taskOverlayDisabled_when_usingDesktopTask() =
+        testScope.runTest {
+            sut = createTaskViewModel(TaskViewType.DESKTOP)
+            sut.bind(TASK_MODEL_1.id)
+            recentsViewData.overlayEnabled.value = true
+            recentsViewData.settledFullyVisibleTaskIds.value = setOf(1)
+
+            assertThat(sut.state.first().taskOverlayEnabled).isFalse()
+        }
+
+    @Test
+    fun taskOverlayDisabled_when_OverlayIsEnabledForInvisibleTask() =
+        testScope.runTest {
+            sut.bind(TASK_MODEL_1.id)
+            recentsViewData.overlayEnabled.value = true
+            recentsViewData.settledFullyVisibleTaskIds.value = setOf(2)
+
+            assertThat(sut.state.first().taskOverlayEnabled).isFalse()
+        }
+
+    @Test
+    fun taskOverlayDisabled_when_OverlayIsDisabledForVisibleTask() =
+        testScope.runTest {
+            sut.bind(TASK_MODEL_1.id)
+            recentsViewData.overlayEnabled.value = false
+            recentsViewData.settledFullyVisibleTaskIds.value = setOf(1)
+
+            assertThat(sut.state.first().taskOverlayEnabled).isFalse()
+        }
+
+    @Test
     fun shouldShowSplash_calls_useCase() {
         sut.isThumbnailValid(null, 0, 0)
         verify(isThumbnailValidUseCase).invoke(anyOrNull(), anyInt(), anyInt())
@@ -256,6 +306,17 @@
             isLocked = isLocked,
         )
 
+    private fun createTaskViewModel(taskViewType: TaskViewType) =
+        TaskViewModel(
+            taskViewType = taskViewType,
+            recentsViewData = recentsViewData,
+            getTaskUseCase = getTaskUseCase,
+            getSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase(),
+            isThumbnailValidUseCase = isThumbnailValidUseCase,
+            getThumbnailPositionUseCase = getThumbnailPositionUseCase,
+            dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
+        )
+
     private companion object {
         const val INVALID_TASK_ID = -1
         const val FLAGS_APPEARANCE_LIGHT_THEME = FLAG_LIGHT_STATUS or FLAG_LIGHT_NAV