diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index 4771d1e..bb5728d 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -76,7 +76,7 @@
         mSecondaryTask = secondary;
         mTaskIdContainer[1] = secondary.key.id;
         mTaskIdAttributeContainer[1] = new TaskIdAttributeContainer(secondary, mSnapshotView2,
-                mIconView2, STAGE_POSITION_BOTTOM_OR_RIGHT);
+                STAGE_POSITION_BOTTOM_OR_RIGHT);
         mTaskIdAttributeContainer[0].setStagePosition(STAGE_POSITION_TOP_OR_LEFT);
         mSnapshotView2.bind(secondary);
         mSplitBoundsConfig = splitBoundsConfig;
@@ -117,6 +117,14 @@
         }
     }
 
+    protected boolean showTaskMenuWithContainer(IconView iconView) {
+        if (mActivity.getDeviceProfile().overviewShowAsGrid) {
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
+        } else {
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        }
+    }
+
     public void updateSplitBoundsConfig(StagedSplitBounds stagedSplitBounds) {
         mSplitBoundsConfig = stagedSplitBounds;
         invalidate();
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
index d73b731..179fd68 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
@@ -80,6 +80,7 @@
 
     private val menuWidth = context.resources.getDimensionPixelSize(R.dimen.task_menu_width_grid)
 
+    private lateinit var taskView: TaskView
     private lateinit var optionLayout: LinearLayout
     private lateinit var taskContainer: TaskIdAttributeContainer
 
@@ -93,7 +94,7 @@
     override fun isOfType(type: Int): Boolean = type and TYPE_TASK_MENU != 0
 
     override fun getTargetObjectLocation(outPos: Rect?) {
-        popupContainer.getDescendantRectRelativeToSelf(taskContainer.iconView, outPos)
+        popupContainer.getDescendantRectRelativeToSelf(taskView.iconView, outPos)
     }
 
     override fun onControllerInterceptTouchEvent(ev: MotionEvent?): Boolean {
@@ -116,6 +117,7 @@
             return false
         }
 
+        taskView = taskContainer.taskView
         this.taskContainer = taskContainer
         if (!populateMenu()) return false
         show()
@@ -135,8 +137,7 @@
     private fun addMenuOptions() {
         // Add the options
         TaskOverlayFactory
-            .getEnabledShortcuts(taskContainer.taskView, mActivityContext.deviceProfile,
-                    taskContainer)
+            .getEnabledShortcuts(taskView, mActivityContext.deviceProfile, taskContainer)
             .forEach { this.addMenuOption(it) }
 
         // Add the spaces between items
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 1ff2a88..67128f0 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -539,7 +539,7 @@
         mTask = task;
         mTaskIdContainer[0] = mTask.key.id;
         mTaskIdAttributeContainer[0] = new TaskIdAttributeContainer(task, mSnapshotView,
-                mIconView, STAGE_POSITION_UNDEFINED);
+                STAGE_POSITION_UNDEFINED);
         mSnapshotView.bind(task);
         setOrientationState(orientedState);
     }
@@ -837,12 +837,10 @@
     }
 
     protected boolean showTaskMenuWithContainer(IconView iconView) {
-        TaskIdAttributeContainer menuContainer =
-                mTaskIdAttributeContainer[iconView == mIconView ? 0 : 1];
         if (mActivity.getDeviceProfile().overviewShowAsGrid) {
-            return TaskMenuViewWithArrow.Companion.showForTask(menuContainer);
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
         } else {
-            return TaskMenuView.showForTask(menuContainer);
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
         }
     }
 
@@ -1564,15 +1562,13 @@
     public class TaskIdAttributeContainer {
         private final TaskThumbnailView mThumbnailView;
         private final Task mTask;
-        private final IconView mIconView;
         /** Defaults to STAGE_POSITION_UNDEFINED if in not a split screen task view */
         private @SplitConfigurationOptions.StagePosition int mStagePosition;
 
         public TaskIdAttributeContainer(Task task, TaskThumbnailView thumbnailView,
-                IconView iconView, int stagePosition) {
+                int stagePosition) {
             this.mTask = task;
             this.mThumbnailView = thumbnailView;
-            this.mIconView = iconView;
             this.mStagePosition = stagePosition;
         }
 
@@ -1592,10 +1588,6 @@
             return TaskView.this;
         }
 
-        public IconView getIconView() {
-            return mIconView;
-        }
-
         public int getStagePosition() {
             return mStagePosition;
         }