diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
index 4c678a2bf29be..28fcdde08e259 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
@@ -696,9 +696,10 @@ public class DefaultTransitionHandler implements Transitions.TransitionHandler {
 
         if (a != null) {
             if (!a.isInitialized()) {
-                final int width = endBounds.width();
-                final int height = endBounds.height();
-                a.initialize(width, height, width, height);
+                final Rect animationRange = TransitionUtil.isClosingType(changeMode)
+                        ? change.getStartAbsBounds() : change.getEndAbsBounds();
+                a.initialize(animationRange.width(), animationRange.height(),
+                        endBounds.width(), endBounds.height());
             }
             a.restrictDuration(MAX_ANIMATION_DURATION);
             a.scaleCurrentDuration(mTransitionAnimationScaleSetting);
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index bb6f8056acda7..7104739068e4b 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -3290,6 +3290,12 @@ class Task extends TaskFragment {
             scheduleAnimation();
         }
 
+        // Let organizer manage task visibility for shell transition. So don't change it's
+        // visibility during collecting.
+        if (mTransitionController.isCollecting() && mCreatedByOrganizer) {
+            return;
+        }
+
         // We intend to let organizer manage task visibility but it doesn't
         // have enough information until we finish shell transitions.
         // In the mean time we do an easy fix here.
@@ -5687,17 +5693,28 @@ class Task extends TaskFragment {
     }
 
     private boolean moveTaskToBackInner(@NonNull Task task) {
-        moveToBack("moveTaskToBackInner", task);
-
-        if (inPinnedWindowingMode()) {
-            mTaskSupervisor.removeRootTask(this);
-            return true;
+        if (mTransitionController.isShellTransitionsEnabled()) {
+            // Preventing from update surface position for WindowState if configuration changed,
+            // because the position is depends on WindowFrame, so update the position before
+            // relayout will only update it to "old" position.
+            mAtmService.deferWindowLayout();
         }
+        try {
+            moveToBack("moveTaskToBackInner", task);
 
-        mRootWindowContainer.ensureVisibilityAndConfig(null /* starting */,
-                mDisplayContent.mDisplayId, false /* markFrozenIfConfigChanged */,
-                false /* deferResume */);
+            if (inPinnedWindowingMode()) {
+                mTaskSupervisor.removeRootTask(this);
+                return true;
+            }
 
+            mRootWindowContainer.ensureVisibilityAndConfig(null /* starting */,
+                    mDisplayContent.mDisplayId, false /* markFrozenIfConfigChanged */,
+                    false /* deferResume */);
+        } finally {
+            if (mTransitionController.isShellTransitionsEnabled()) {
+                mAtmService.continueWindowLayout();
+            }
+        }
         ActivityRecord topActivity = getDisplayArea().topRunningActivity();
         Task topRootTask = topActivity.getRootTask();
         if (topRootTask != null && topRootTask != this && topActivity.isState(RESUMED)) {
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index b938b9e4ec51e..d9b13b07470f7 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -2245,11 +2245,17 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
             WindowContainer<?> ancestor = findCommonAncestor(sortedTargets, wc);
 
             // Make leash based on highest (z-order) direct child of ancestor with a participant.
+            // Check whether the ancestor is belonged to last parent, shouldn't happen.
+            final boolean hasReparent = !wc.isDescendantOf(ancestor);
             WindowContainer leashReference = wc;
-            while (leashReference.getParent() != ancestor) {
-                leashReference = leashReference.getParent();
+            if (hasReparent) {
+                Slog.e(TAG, "Did not find common ancestor! Ancestor= " + ancestor
+                        + " target= " + wc);
+            } else {
+                while (leashReference.getParent() != ancestor) {
+                    leashReference = leashReference.getParent();
+                }
             }
-
             final SurfaceControl rootLeash = leashReference.makeAnimationLeash().setName(
                     "Transition Root: " + leashReference.getName()).build();
             rootLeash.setUnreleasedWarningCallSite("Transition.calculateTransitionRoots");
@@ -2454,6 +2460,20 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
                 // Skip the non-app window or windows on a different display
                 continue;
             }
+            // Re-initiate the last parent as the initial ancestor instead of the top target.
+            // When move a leaf task from organized task to display area, try to keep the transition
+            // root be the original organized task for close transition animation.
+            // Otherwise, shell will use wrong root layer to play animation.
+            // Note: Since the target is sorted, so only need to do this at the lowest target.
+            if (change.mStartParent != null && wc.getParent() != null
+                    && change.mStartParent.isAttached() && wc.getParent() != change.mStartParent
+                    && i == targets.size() - 1) {
+                final int transitionMode = change.getTransitMode(wc);
+                if (transitionMode == TRANSIT_CLOSE || transitionMode == TRANSIT_TO_BACK) {
+                    ancestor = change.mStartParent;
+                    continue;
+                }
+            }
             while (!wc.isDescendantOf(ancestor)) {
                 ancestor = ancestor.getParent();
             }