diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index bb5728d..4771d1e 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,
-                STAGE_POSITION_BOTTOM_OR_RIGHT);
+                mIconView2, STAGE_POSITION_BOTTOM_OR_RIGHT);
         mTaskIdAttributeContainer[0].setStagePosition(STAGE_POSITION_TOP_OR_LEFT);
         mSnapshotView2.bind(secondary);
         mSplitBoundsConfig = splitBoundsConfig;
@@ -117,14 +117,6 @@
         }
     }
 
-    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 179fd68..cd1691b 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
@@ -94,7 +94,7 @@
     override fun isOfType(type: Int): Boolean = type and TYPE_TASK_MENU != 0
 
     override fun getTargetObjectLocation(outPos: Rect?) {
-        popupContainer.getDescendantRectRelativeToSelf(taskView.iconView, outPos)
+        popupContainer.getDescendantRectRelativeToSelf(taskContainer.iconView, outPos)
     }
 
     override fun onControllerInterceptTouchEvent(ev: MotionEvent?): Boolean {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 67128f0..1ff2a88 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,
-                STAGE_POSITION_UNDEFINED);
+                mIconView, STAGE_POSITION_UNDEFINED);
         mSnapshotView.bind(task);
         setOrientationState(orientedState);
     }
@@ -837,10 +837,12 @@
     }
 
     protected boolean showTaskMenuWithContainer(IconView iconView) {
+        TaskIdAttributeContainer menuContainer =
+                mTaskIdAttributeContainer[iconView == mIconView ? 0 : 1];
         if (mActivity.getDeviceProfile().overviewShowAsGrid) {
-            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
+            return TaskMenuViewWithArrow.Companion.showForTask(menuContainer);
         } else {
-            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+            return TaskMenuView.showForTask(menuContainer);
         }
     }
 
@@ -1562,13 +1564,15 @@
     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,
-                int stagePosition) {
+                IconView iconView, int stagePosition) {
             this.mTask = task;
             this.mThumbnailView = thumbnailView;
+            this.mIconView = iconView;
             this.mStagePosition = stagePosition;
         }
 
@@ -1588,6 +1592,10 @@
             return TaskView.this;
         }
 
+        public IconView getIconView() {
+            return mIconView;
+        }
+
         public int getStagePosition() {
             return mStagePosition;
         }