diff --git a/packages/SystemUI/src/com/android/systemui/recents/model/TaskStack.java b/packages/SystemUI/src/com/android/systemui/recents/model/TaskStack.java
index 50e28ca2205d9..26200d0fac477 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/model/TaskStack.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/model/TaskStack.java
@@ -228,8 +228,8 @@ public class TaskStack {
         /**
          * Notifies when a task has been removed from the stack.
          */
-        void onStackTaskRemoved(TaskStack stack, Task removedTask, boolean wasFrontMostTask,
-            Task newFrontMostTask, AnimationProps animation, boolean fromDockGesture);
+        void onStackTaskRemoved(TaskStack stack, Task removedTask, Task newFrontMostTask,
+                AnimationProps animation, boolean fromDockGesture);
 
         /**
          * Notifies when all tasks have been removed from the stack.
@@ -609,12 +609,11 @@ public class TaskStack {
      */
     public void removeTask(Task t, AnimationProps animation, boolean fromDockGesture) {
         if (mStackTaskList.contains(t)) {
-            boolean wasFrontMostTask = (getStackFrontMostTask(false /* includeFreeform */) == t);
             removeTaskImpl(mStackTaskList, t);
             Task newFrontMostTask = getStackFrontMostTask(false  /* includeFreeform */);
             if (mCb != null) {
                 // Notify that a task has been removed
-                mCb.onStackTaskRemoved(this, t, wasFrontMostTask, newFrontMostTask, animation,
+                mCb.onStackTaskRemoved(this, t, newFrontMostTask, animation,
                         fromDockGesture);
             }
         }
@@ -648,6 +647,7 @@ public class TaskStack {
         ArrayMap<Task.TaskKey, Task> currentTasksMap = createTaskKeyMapFromList(mRawTaskList);
         ArrayMap<Task.TaskKey, Task> newTasksMap = createTaskKeyMapFromList(tasks);
         ArrayList<Task> addedTasks = new ArrayList<>();
+        ArrayList<Task> removedTasks = new ArrayList<>();
         ArrayList<Task> allTasks = new ArrayList<>();
 
         // Disable notifications if there are no callbacks
@@ -661,11 +661,7 @@ public class TaskStack {
             Task task = mRawTaskList.get(i);
             if (!newTasksMap.containsKey(task.key)) {
                 if (notifyStackChanges) {
-                    // If we are notifying, then remove the task now, otherwise the raw task list
-                    // will be reset at the end of this method
-                    removeTask(task, AnimationProps.IMMEDIATE, false /* fromDockGesture */);
-                    mCb.onStackTaskRemoved(this, task, i == (taskCount - 1), null,
-                            AnimationProps.IMMEDIATE, false /* fromDockGesture */);
+                    removedTasks.add(task);
                 }
             }
             task.setGroup(null);
@@ -699,6 +695,14 @@ public class TaskStack {
         // Update the affiliated groupings
         createAffiliatedGroupings(context);
 
+        // Only callback for the removed tasks after the stack has updated
+        int removedTaskCount = removedTasks.size();
+        Task newFrontMostTask = getStackFrontMostTask(false);
+        for (int i = 0; i < removedTaskCount; i++) {
+            mCb.onStackTaskRemoved(this, removedTasks.get(i), newFrontMostTask,
+                    AnimationProps.IMMEDIATE, false /* fromDockGesture */);
+        }
+
         // Only callback for the newly added tasks after this stack has been updated
         int addedTaskCount = addedTasks.size();
         for (int i = 0; i < addedTaskCount; i++) {
diff --git a/packages/SystemUI/src/com/android/systemui/recents/tv/views/TaskStackHorizontalGridView.java b/packages/SystemUI/src/com/android/systemui/recents/tv/views/TaskStackHorizontalGridView.java
index 9f52abdad287f..3fb339e8f98a2 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/tv/views/TaskStackHorizontalGridView.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/tv/views/TaskStackHorizontalGridView.java
@@ -181,8 +181,8 @@ public class TaskStackHorizontalGridView extends HorizontalGridView implements T
     }
 
     @Override
-    public void onStackTaskRemoved(TaskStack stack, Task removedTask, boolean wasFrontMostTask,
-            Task newFrontMostTask, AnimationProps animation, boolean fromDockGesture) {
+    public void onStackTaskRemoved(TaskStack stack, Task removedTask, Task newFrontMostTask,
+            AnimationProps animation, boolean fromDockGesture) {
         ((TaskStackHorizontalViewAdapter) getAdapter()).removeTask(removedTask);
         if (mFocusedTask == removedTask) {
             resetFocusedTask(removedTask);
diff --git a/packages/SystemUI/src/com/android/systemui/recents/views/TaskStackView.java b/packages/SystemUI/src/com/android/systemui/recents/views/TaskStackView.java
index 1a197b69c1bb8..0d9f77c5d2faa 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/views/TaskStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/views/TaskStackView.java
@@ -1400,8 +1400,8 @@ public class TaskStackView extends FrameLayout implements TaskStack.TaskStackCal
      * We expect that the {@link TaskView} associated with the removed task is already hidden.
      */
     @Override
-    public void onStackTaskRemoved(TaskStack stack, Task removedTask, boolean wasFrontMostTask,
-            Task newFrontMostTask, AnimationProps animation, boolean fromDockGesture) {
+    public void onStackTaskRemoved(TaskStack stack, Task removedTask, Task newFrontMostTask,
+            AnimationProps animation, boolean fromDockGesture) {
         if (mFocusedTask == removedTask) {
             resetFocusedTask(removedTask);
         }