diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index 793bad86d873d..48487bc4a3d60 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -129,11 +129,11 @@ public class FreeformTaskListener implements ShellTaskOrganizer.TaskListener,
     @Override
     public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
         final State state = mTasks.get(taskInfo.taskId);
-        state.mTaskInfo = taskInfo;
+
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TASK_ORG, "Freeform Task Info Changed: #%d",
                 taskInfo.taskId);
-        mWindowDecorationViewModel.onTaskInfoChanged(state.mTaskInfo);
-
+        mWindowDecorationViewModel.onTaskInfoChanged(taskInfo);
+        state.mTaskInfo = taskInfo;
         if (DesktopModeStatus.isAnyEnabled()) {
             mDesktopModeTaskRepository.ifPresent(repository -> {
                 if (taskInfo.isVisible) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
index 6623f5ca84eef..b9caf62012a61 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
@@ -123,10 +123,11 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
     public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
         final State state = mTasks.get(taskInfo.taskId);
         final Point oldPositionInParent = state.mTaskInfo.positionInParent;
-        state.mTaskInfo = taskInfo;
+
         if (mWindowDecorViewModelOptional.isPresent()) {
-            mWindowDecorViewModelOptional.get().onTaskInfoChanged(state.mTaskInfo);
+            mWindowDecorViewModelOptional.get().onTaskInfoChanged(taskInfo);
         }
+        state.mTaskInfo = taskInfo;
         if (Transitions.ENABLE_SHELL_TRANSITIONS) return;
         updateRecentsForVisibleFullscreenTask(taskInfo);
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index b4301577ac101..00aab67333692 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -153,12 +153,11 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
     @Override
     public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
         final DesktopModeWindowDecoration decoration = mWindowDecorByTaskId.get(taskInfo.taskId);
-
         if (decoration == null) return;
+        final RunningTaskInfo oldTaskInfo = decoration.mTaskInfo;
 
-        int oldDisplayId = decoration.mDisplay.getDisplayId();
-        if (taskInfo.displayId != oldDisplayId) {
-            removeTaskFromEventReceiver(oldDisplayId);
+        if (taskInfo.displayId != oldTaskInfo.displayId) {
+            removeTaskFromEventReceiver(oldTaskInfo.displayId);
             incrementEventReceiverTasks(taskInfo.displayId);
         }