Merge "[PB] Migrate visibility change to transition." into main
diff --git a/core/java/android/view/WindowManager.java b/core/java/android/view/WindowManager.java
index 14978ed..85d4ec0 100644
--- a/core/java/android/view/WindowManager.java
+++ b/core/java/android/view/WindowManager.java
@@ -478,6 +478,12 @@
      */
     int TRANSIT_SLEEP = 12;
     /**
+     * An Activity was going to be visible from back navigation.
+     * @hide
+     */
+    int TRANSIT_PREPARE_BACK_NAVIGATION = 13;
+
+    /**
      * The first slot for custom transition types. Callers (like Shell) can make use of custom
      * transition types for dealing with special cases. These types are effectively ignored by
      * Core and will just be passed along as part of TransitionInfo objects. An example is
@@ -505,6 +511,7 @@
             TRANSIT_PIP,
             TRANSIT_WAKE,
             TRANSIT_SLEEP,
+            TRANSIT_PREPARE_BACK_NAVIGATION,
             TRANSIT_FIRST_CUSTOM
     })
     @Retention(RetentionPolicy.SOURCE)
@@ -1918,6 +1925,7 @@
             case TRANSIT_PIP: return "PIP";
             case TRANSIT_WAKE: return "WAKE";
             case TRANSIT_SLEEP: return "SLEEP";
+            case TRANSIT_PREPARE_BACK_NAVIGATION: return "PREDICTIVE_BACK";
             case TRANSIT_FIRST_CUSTOM: return "FIRST_CUSTOM";
             default:
                 if (type > TRANSIT_FIRST_CUSTOM) {
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 80101af..80a0102 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -209,3 +209,13 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+flag {
+  name: "migrate_predictive_back_transition"
+  namespace: "windowing_frontend"
+  description: "Create transition when visibility change from predictive back"
+  bug: "347168362"
+  is_fixed_read_only: true
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
index a9fdea3..f14f419 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
@@ -16,7 +16,12 @@
 
 package com.android.wm.shell.back;
 
+import static android.view.RemoteAnimationTarget.MODE_CLOSING;
+import static android.view.RemoteAnimationTarget.MODE_OPENING;
+import static android.window.TransitionInfo.FLAG_BACK_GESTURE_ANIMATED;
+
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_PREDICTIVE_BACK_HOME;
+import static com.android.window.flags.Flags.migratePredictiveBackTransition;
 import static com.android.window.flags.Flags.predictiveBackSystemAnims;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_BACK_PREVIEW;
 import static com.android.wm.shell.sysui.ShellSharedConstants.KEY_EXTRA_SHELL_BACK_ANIMATION;
@@ -30,11 +35,13 @@
 import android.content.Context;
 import android.content.res.Configuration;
 import android.database.ContentObserver;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.input.InputManager;
 import android.net.Uri;
 import android.os.Bundle;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.os.SystemClock;
@@ -48,6 +55,7 @@
 import android.view.KeyEvent;
 import android.view.MotionEvent;
 import android.view.RemoteAnimationTarget;
+import android.view.SurfaceControl;
 import android.view.WindowManager;
 import android.window.BackAnimationAdapter;
 import android.window.BackEvent;
@@ -57,6 +65,9 @@
 import android.window.IBackAnimationFinishedCallback;
 import android.window.IBackAnimationRunner;
 import android.window.IOnBackInvokedCallback;
+import android.window.TransitionInfo;
+import android.window.TransitionRequestInfo;
+import android.window.WindowContainerTransaction;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.ProtoLog;
@@ -66,12 +77,14 @@
 import com.android.wm.shell.common.ExternalInterfaceBinder;
 import com.android.wm.shell.common.RemoteCallable;
 import com.android.wm.shell.common.ShellExecutor;
+import com.android.wm.shell.shared.TransitionUtil;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
 import com.android.wm.shell.shared.annotations.ShellMainThread;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.transition.Transitions;
 
 import java.io.PrintWriter;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -101,6 +114,7 @@
 
     /** Tracks if an uninterruptible animation is in progress */
     private boolean mPostCommitAnimationInProgress = false;
+    private boolean mRealCallbackInvoked = false;
 
     /** Tracks if we should start the back gesture on the next motion move event */
     private boolean mShouldStartOnNextMoveEvent = false;
@@ -123,6 +137,8 @@
     private final ShellExecutor mShellExecutor;
     private final Handler mBgHandler;
     private final WindowManager mWindowManager;
+    private final Transitions mTransitions;
+    private final BackTransitionHandler mBackTransitionHandler;
     @VisibleForTesting
     final Rect mTouchableArea = new Rect();
 
@@ -190,7 +206,8 @@
             Context context,
             @NonNull BackAnimationBackground backAnimationBackground,
             ShellBackAnimationRegistry shellBackAnimationRegistry,
-            ShellCommandHandler shellCommandHandler) {
+            ShellCommandHandler shellCommandHandler,
+            Transitions transitions) {
         this(
                 shellInit,
                 shellController,
@@ -201,7 +218,8 @@
                 context.getContentResolver(),
                 backAnimationBackground,
                 shellBackAnimationRegistry,
-                shellCommandHandler);
+                shellCommandHandler,
+                transitions);
     }
 
     @VisibleForTesting
@@ -215,7 +233,8 @@
             ContentResolver contentResolver,
             @NonNull BackAnimationBackground backAnimationBackground,
             ShellBackAnimationRegistry shellBackAnimationRegistry,
-            ShellCommandHandler shellCommandHandler) {
+            ShellCommandHandler shellCommandHandler,
+            Transitions transitions) {
         mShellController = shellController;
         mShellExecutor = shellExecutor;
         mActivityTaskManager = activityTaskManager;
@@ -230,6 +249,9 @@
         mLatencyTracker = LatencyTracker.getInstance(mContext);
         mShellCommandHandler = shellCommandHandler;
         mWindowManager = context.getSystemService(WindowManager.class);
+        mTransitions = transitions;
+        mBackTransitionHandler = new BackTransitionHandler();
+        mTransitions.addHandler(mBackTransitionHandler);
         updateTouchableArea();
     }
 
@@ -730,7 +752,7 @@
             mBackAnimationFinishedCallback = null;
         }
 
-        if (mBackNavigationInfo != null) {
+        if (mBackNavigationInfo != null && !mRealCallbackInvoked) {
             final IOnBackInvokedCallback callback = mBackNavigationInfo.getOnBackInvokedCallback();
             if (touchTracker.getTriggerBack()) {
                 dispatchOnBackInvoked(callback);
@@ -738,6 +760,7 @@
                 tryDispatchOnBackCancelled(callback);
             }
         }
+        mRealCallbackInvoked = false;
         finishBackNavigation(touchTracker.getTriggerBack());
     }
 
@@ -815,14 +838,38 @@
 
         // The next callback should be {@link #onBackAnimationFinished}.
         if (mCurrentTracker.getTriggerBack()) {
-            // notify gesture finished
-            mBackNavigationInfo.onBackGestureFinished(true);
+            if (migratePredictiveBackTransition()) {
+                // notify core gesture is commit
+                if (shouldTriggerCloseTransition()) {
+                    mBackTransitionHandler.mCloseTransitionRequested = true;
+                    final IOnBackInvokedCallback callback =
+                            mBackNavigationInfo.getOnBackInvokedCallback();
+                    // invoked client side onBackInvoked
+                    dispatchOnBackInvoked(callback);
+                    mRealCallbackInvoked = true;
+                }
+            } else {
+                // notify gesture finished
+                mBackNavigationInfo.onBackGestureFinished(true);
+            }
+
+            // start post animation
             dispatchOnBackInvoked(mActiveCallback);
         } else {
             tryDispatchOnBackCancelled(mActiveCallback);
         }
     }
 
+    // Close window won't create any transition
+    private boolean shouldTriggerCloseTransition() {
+        if (mBackNavigationInfo == null) {
+            return false;
+        }
+        int type = mBackNavigationInfo.getType();
+        return type == BackNavigationInfo.TYPE_RETURN_TO_HOME
+                || type == BackNavigationInfo.TYPE_CROSS_TASK
+                || type == BackNavigationInfo.TYPE_CROSS_ACTIVITY;
+    }
     /**
      * Called when the post commit animation is completed or timeout.
      * This will trigger the real {@link IOnBackInvokedCallback} behavior.
@@ -857,6 +904,7 @@
                     "mCurrentBackGestureInfo was null when back animation finished");
         }
         resetTouchTracker();
+        mBackTransitionHandler.onAnimationFinished();
     }
 
     /**
@@ -1016,11 +1064,13 @@
                                     endLatencyTracking();
                                     if (!validateAnimationTargets(apps)) {
                                         Log.e(TAG, "Invalid animation targets!");
+                                        mBackTransitionHandler.consumeQueuedTransitionIfNeeded();
                                         return;
                                     }
                                     mBackAnimationFinishedCallback = finishedCallback;
                                     mApps = apps;
                                     startSystemAnimation();
+                                    mBackTransitionHandler.consumeQueuedTransitionIfNeeded();
 
                                     // Dispatch the first progress after animation start for
                                     // smoothing the initial animation, instead of waiting for next
@@ -1041,6 +1091,7 @@
                     public void onAnimationCancelled() {
                         mShellExecutor.execute(
                                 () -> {
+                                    mBackTransitionHandler.consumeQueuedTransitionIfNeeded();
                                     if (!mShellBackAnimationRegistry.cancel(
                                             mBackNavigationInfo != null
                                                     ? mBackNavigationInfo.getType()
@@ -1073,4 +1124,249 @@
         mQueuedTracker.dump(pw, prefix + "    ");
     }
 
+    class BackTransitionHandler implements Transitions.TransitionHandler {
+
+        Runnable mOnAnimationFinishCallback;
+        boolean mCloseTransitionRequested;
+        boolean mOpeningRunning;
+        SurfaceControl.Transaction mFinishOpenTransaction;
+        Transitions.TransitionFinishCallback mFinishOpenTransitionCallback;
+        QueuedTransition mQueuedTransition = null;
+        void onAnimationFinished() {
+            if (!mCloseTransitionRequested) {
+                applyFinishOpenTransition();
+            }
+            if (mOnAnimationFinishCallback != null) {
+                mOnAnimationFinishCallback.run();
+                mOnAnimationFinishCallback = null;
+            }
+        }
+
+        void consumeQueuedTransitionIfNeeded() {
+            if (mQueuedTransition != null) {
+                mQueuedTransition.consume();
+                mQueuedTransition = null;
+            }
+        }
+
+        private void applyFinishOpenTransition() {
+            if (mFinishOpenTransaction != null) {
+                mFinishOpenTransaction.apply();
+                mFinishOpenTransaction = null;
+            }
+            if (mFinishOpenTransitionCallback != null) {
+                mFinishOpenTransitionCallback.onTransitionFinished(null);
+                mFinishOpenTransitionCallback = null;
+            }
+            mOpeningRunning = false;
+        }
+
+        private void applyAndFinish(@NonNull SurfaceControl.Transaction st,
+                @NonNull SurfaceControl.Transaction ft,
+                @NonNull Transitions.TransitionFinishCallback finishCallback) {
+            applyFinishOpenTransition();
+            st.apply();
+            ft.apply();
+            finishCallback.onTransitionFinished(null);
+            mCloseTransitionRequested = false;
+        }
+        @Override
+        public boolean startAnimation(@NonNull IBinder transition,
+                @NonNull TransitionInfo info,
+                @NonNull SurfaceControl.Transaction st,
+                @NonNull SurfaceControl.Transaction ft,
+                @NonNull Transitions.TransitionFinishCallback finishCallback) {
+            // Both mShellExecutor and Transitions#mMainExecutor are ShellMainThread, so we don't
+            // need to post to ShellExecutor when called.
+            if (info.getType() != WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION
+                    && !isGestureBackTransition(info)) {
+                return false;
+            }
+            if (mApps == null || mApps.length == 0) {
+                if (mBackNavigationInfo != null && mShellBackAnimationRegistry
+                        .isWaitingAnimation(mBackNavigationInfo.getType())) {
+                    // Waiting for animation? Queue update to wait for animation start.
+                    consumeQueuedTransitionIfNeeded();
+                    mQueuedTransition = new QueuedTransition(info, st, ft, finishCallback);
+                } else {
+                    // animation was done, consume directly
+                    applyAndFinish(st, ft, finishCallback);
+                }
+                return true;
+            }
+
+            if (handlePrepareTransition(info, st, ft, finishCallback)) {
+                return true;
+            }
+            return handleCloseTransition(info, st, ft, finishCallback);
+        }
+
+        @Override
+        public void mergeAnimation(@NonNull IBinder transition, @NonNull TransitionInfo info,
+                @NonNull SurfaceControl.Transaction t, @NonNull IBinder mergeTarget,
+                @NonNull Transitions.TransitionFinishCallback finishCallback) {
+            if (!isGestureBackTransition(info)) {
+                if (mOpeningRunning) {
+                    applyFinishOpenTransition();
+                }
+                if (mQueuedTransition != null) {
+                    consumeQueuedTransitionIfNeeded();
+                }
+                return;
+            }
+            // Handle the commit transition if this handler is running the open transition.
+            finishCallback.onTransitionFinished(null);
+            if (mCloseTransitionRequested) {
+                if (mApps == null || mApps.length == 0) {
+                    if (mQueuedTransition == null) {
+                        // animation was done
+                        applyFinishOpenTransition();
+                        mCloseTransitionRequested = false;
+                    } // else, let queued transition to play
+                } else {
+                    // we are animating, wait until animation finish
+                    mOnAnimationFinishCallback = () -> {
+                        applyFinishOpenTransition();
+                        mCloseTransitionRequested = false;
+                    };
+                }
+            }
+        }
+
+        /**
+         * Check whether this transition is prepare for predictive back animation, which could
+         * happen when core make an activity become visible.
+         */
+        private boolean handlePrepareTransition(
+                @NonNull TransitionInfo info,
+                @NonNull SurfaceControl.Transaction st,
+                @NonNull SurfaceControl.Transaction ft,
+                @NonNull Transitions.TransitionFinishCallback finishCallback) {
+            if (info.getType() != WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION) {
+                return false;
+            }
+
+            SurfaceControl openingLeash = null;
+            for (int i = mApps.length - 1; i >= 0; --i) {
+                if (mApps[i].mode == MODE_OPENING) {
+                    openingLeash = mApps[i].leash;
+                }
+            }
+            if (openingLeash != null) {
+                for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+                    final TransitionInfo.Change c = info.getChanges().get(i);
+                    if (TransitionUtil.isOpeningMode(c.getMode())) {
+                        final Point offset = c.getEndRelOffset();
+                        st.setPosition(c.getLeash(), offset.x, offset.y);
+                        st.reparent(c.getLeash(), openingLeash);
+                    }
+                }
+            }
+            st.apply();
+            mFinishOpenTransaction = ft;
+            mFinishOpenTransitionCallback = finishCallback;
+            mOpeningRunning = true;
+            return true;
+        }
+
+        private boolean isGestureBackTransition(@NonNull TransitionInfo info) {
+            for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+                final TransitionInfo.Change c = info.getChanges().get(i);
+                if (c.hasFlags(FLAG_BACK_GESTURE_ANIMATED)
+                        && (TransitionUtil.isOpeningMode(c.getMode())
+                        || TransitionUtil.isClosingMode(c.getMode()))) {
+                    return true;
+                }
+            }
+            return false;
+        }
+        /**
+         * Check whether this transition is triggered from back gesture commitment.
+         * Reparent the transition targets to animation leashes, so the animation won't be broken.
+         */
+        private boolean handleCloseTransition(@NonNull TransitionInfo info,
+                @NonNull SurfaceControl.Transaction st,
+                @NonNull SurfaceControl.Transaction ft,
+                @NonNull Transitions.TransitionFinishCallback finishCallback) {
+            SurfaceControl openingLeash = null;
+            SurfaceControl closingLeash = null;
+            for (int i = mApps.length - 1; i >= 0; --i) {
+                if (mApps[i].mode == MODE_OPENING) {
+                    openingLeash = mApps[i].leash;
+                }
+                if (mApps[i].mode == MODE_CLOSING) {
+                    closingLeash = mApps[i].leash;
+                }
+            }
+            if (openingLeash != null && closingLeash != null) {
+                for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+                    final TransitionInfo.Change c = info.getChanges().get(i);
+                    if (TransitionUtil.isOpeningMode(c.getMode())) {
+                        final Point offset = c.getEndRelOffset();
+                        st.setPosition(c.getLeash(), offset.x, offset.y);
+                        st.reparent(c.getLeash(), openingLeash);
+                    } else if (TransitionUtil.isClosingMode(c.getMode())) {
+                        st.reparent(c.getLeash(), closingLeash);
+                    }
+                }
+            }
+            st.apply();
+            // mApps must exists
+            mOnAnimationFinishCallback = () -> {
+                ft.apply();
+                finishCallback.onTransitionFinished(null);
+                mCloseTransitionRequested = false;
+            };
+            return true;
+        }
+
+        @Nullable
+        @Override
+        public WindowContainerTransaction handleRequest(
+                @NonNull IBinder transition,
+                @NonNull TransitionRequestInfo request) {
+            if (request.getType() == WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION) {
+                return new WindowContainerTransaction();
+            }
+            if (TransitionUtil.isClosingType(request.getType()) && mCloseTransitionRequested) {
+                return new WindowContainerTransaction();
+            }
+            return null;
+        }
+
+        class QueuedTransition {
+            final TransitionInfo mInfo;
+            final SurfaceControl.Transaction mSt;
+            final  SurfaceControl.Transaction mFt;
+            final Transitions.TransitionFinishCallback mFinishCallback;
+            QueuedTransition(@NonNull TransitionInfo info,
+                    @NonNull SurfaceControl.Transaction st,
+                    @NonNull SurfaceControl.Transaction ft,
+                    @NonNull Transitions.TransitionFinishCallback finishCallback) {
+                mInfo = info;
+                mSt = st;
+                mFt = ft;
+                mFinishCallback = finishCallback;
+            }
+
+            void consume() {
+                // not animating, consume transition directly
+                if (mApps == null || mApps.length == 0) {
+                    applyAndFinish(mSt, mFt, mFinishCallback);
+                    return;
+                }
+                // we are animating
+                if (handlePrepareTransition(mInfo, mSt, mFt, mFinishCallback)) {
+                    // handle merge transition if any
+                    if (mCloseTransitionRequested) {
+                        mOnAnimationFinishCallback = () -> {
+                            applyFinishOpenTransition();
+                            mCloseTransitionRequested = false;
+                        };
+                    }
+                }
+                handleCloseTransition(mInfo, mSt, mFt, mFinishCallback);
+            }
+        }
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index 2dc6382..717a414 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -375,7 +375,8 @@
             @ShellBackgroundThread Handler backgroundHandler,
             BackAnimationBackground backAnimationBackground,
             Optional<ShellBackAnimationRegistry> shellBackAnimationRegistry,
-            ShellCommandHandler shellCommandHandler) {
+            ShellCommandHandler shellCommandHandler,
+            Transitions transitions) {
         if (BackAnimationController.IS_ENABLED) {
             return shellBackAnimationRegistry.map(
                     (animations) ->
@@ -387,7 +388,8 @@
                                     context,
                                     backAnimationBackground,
                                     animations,
-                                    shellCommandHandler));
+                                    shellCommandHandler,
+                                    transitions));
         }
         return Optional.empty();
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
index 58cf398..c850ff8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
@@ -40,6 +40,7 @@
 
 import static com.android.window.flags.Flags.ensureWallpaperInTransitions;
 import static com.android.systemui.shared.Flags.returnAnimationFrameworkLibrary;
+import static com.android.window.flags.Flags.migratePredictiveBackTransition;
 import static com.android.wm.shell.shared.TransitionUtil.isClosingType;
 import static com.android.wm.shell.shared.TransitionUtil.isOpeningType;
 import static com.android.wm.shell.sysui.ShellSharedConstants.KEY_EXTRA_SHELL_SHELL_TRANSITIONS;
@@ -821,7 +822,8 @@
             }
             // The change has already animated by back gesture, don't need to play transition
             // animation on it.
-            if (change.hasFlags(FLAG_BACK_GESTURE_ANIMATED)) {
+            if (!migratePredictiveBackTransition()
+                    && change.hasFlags(FLAG_BACK_GESTURE_ANIMATED)) {
                 info.getChanges().remove(i);
             }
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
index 57e469d..56fad95 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
@@ -68,6 +68,7 @@
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.sysui.ShellSharedConstants;
+import com.android.wm.shell.transition.Transitions;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -114,6 +115,8 @@
     @Mock
     private ShellCommandHandler mShellCommandHandler;
     @Mock
+    private Transitions mTransitions;
+    @Mock
     private RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
 
     private BackAnimationController mController;
@@ -156,7 +159,8 @@
                         mContentResolver,
                         mAnimationBackground,
                         mShellBackAnimationRegistry,
-                        mShellCommandHandler);
+                        mShellCommandHandler,
+                        mTransitions);
         mShellInit.init();
         mShellExecutor.flushAll();
         mTouchableRegion = new Rect(0, 0, 100, 100);
@@ -316,7 +320,8 @@
                         mContentResolver,
                         mAnimationBackground,
                         mShellBackAnimationRegistry,
-                        mShellCommandHandler);
+                        mShellCommandHandler,
+                        mTransitions);
         shellInit.init();
         registerAnimation(BackNavigationInfo.TYPE_RETURN_TO_HOME);
 
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 8a768c0..c8aa815 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -10686,6 +10686,9 @@
                 return true;
             }
         }
+        if (mAtmService.mBackNavigationController.isStartingSurfaceShown(this)) {
+            return true;
+        }
         if (!super.isSyncFinished(group)) return false;
         if (mDisplayContent != null && mDisplayContent.mUnknownAppVisibilityController
                 .isVisibilityUnknown(this)) {
diff --git a/services/core/java/com/android/server/wm/AppTransitionController.java b/services/core/java/com/android/server/wm/AppTransitionController.java
index 78636a7..5a0cbf3 100644
--- a/services/core/java/com/android/server/wm/AppTransitionController.java
+++ b/services/core/java/com/android/server/wm/AppTransitionController.java
@@ -250,7 +250,7 @@
 
         ArraySet<ActivityRecord> tmpOpenApps = mDisplayContent.mOpeningApps;
         ArraySet<ActivityRecord> tmpCloseApps = mDisplayContent.mClosingApps;
-        if (mDisplayContent.mAtmService.mBackNavigationController.isMonitoringTransition()) {
+        if (mDisplayContent.mAtmService.mBackNavigationController.isMonitoringFinishTransition()) {
             tmpOpenApps = new ArraySet<>(mDisplayContent.mOpeningApps);
             tmpCloseApps = new ArraySet<>(mDisplayContent.mClosingApps);
             if (mDisplayContent.mAtmService.mBackNavigationController
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index 0f8d68b..924f765 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -24,6 +24,7 @@
 import static android.view.WindowManager.TRANSIT_CHANGE;
 import static android.view.WindowManager.TRANSIT_CLOSE;
 import static android.view.WindowManager.TRANSIT_OLD_NONE;
+import static android.view.WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 
 import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_BACK_PREVIEW;
@@ -47,6 +48,7 @@
 import android.os.SystemProperties;
 import android.text.TextUtils;
 import android.util.ArraySet;
+import android.util.Pair;
 import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 import android.view.RemoteAnimationTarget;
@@ -83,11 +85,6 @@
 
     private AnimationHandler mAnimationHandler;
 
-    /**
-     * The transition who match the back navigation targets,
-     * release animation after this transition finish.
-     */
-    private Transition mWaitTransitionFinish;
     private final ArrayList<WindowContainer> mTmpOpenApps = new ArrayList<>();
     private final ArrayList<WindowContainer> mTmpCloseApps = new ArrayList<>();
 
@@ -143,7 +140,7 @@
 
         BackNavigationInfo.Builder infoBuilder = new BackNavigationInfo.Builder();
         synchronized (wmService.mGlobalLock) {
-            if (isMonitoringTransition()) {
+            if (isMonitoringFinishTransition()) {
                 Slog.w(TAG, "Previous animation hasn't finish, status: " + mAnimationHandler);
                 // Don't start any animation for it.
                 return null;
@@ -308,6 +305,7 @@
                     backType = BackNavigationInfo.TYPE_CALLBACK;
                 } else if (prevTask.isActivityTypeHome()) {
                     removedWindowContainer = currentTask;
+                    prevTask = prevTask.getRootTask();
                     backType = BackNavigationInfo.TYPE_RETURN_TO_HOME;
                     final ActivityRecord ar = prevTask.getTopNonFinishingActivity();
                     mShowWallpaper = ar != null && ar.hasWallpaper();
@@ -562,10 +560,15 @@
         return !prevActivities.isEmpty();
     }
 
-    boolean isMonitoringTransition() {
+    boolean isMonitoringFinishTransition() {
         return mAnimationHandler.mComposed || mNavigationMonitor.isMonitorForRemote();
     }
 
+    boolean isMonitoringPrepareTransition(Transition transition) {
+        return mAnimationHandler.mComposed
+                && mAnimationHandler.mOpenAnimAdaptor.mPreparedOpenTransition == transition;
+    }
+
     private void scheduleAnimation(@NonNull AnimationHandler.ScheduleAnimationBuilder builder) {
         mPendingAnimation = builder.build();
         mWindowManagerService.mWindowPlacerLocked.requestTraversal();
@@ -576,7 +579,9 @@
     }
 
     private boolean isWaitBackTransition() {
-        return mAnimationHandler.mComposed && mAnimationHandler.mWaitTransition;
+        // Ignore mWaitTransition while flag is enabled.
+        return mAnimationHandler.mComposed && (Flags.migratePredictiveBackTransition()
+                || mAnimationHandler.mWaitTransition);
     }
 
     boolean isKeyguardOccluded(WindowState focusWindow) {
@@ -626,7 +631,7 @@
      */
     boolean removeIfContainsBackAnimationTargets(ArraySet<ActivityRecord> openApps,
             ArraySet<ActivityRecord> closeApps) {
-        if (!isMonitoringTransition()) {
+        if (!isMonitoringFinishTransition()) {
             return false;
         }
         mTmpCloseApps.addAll(closeApps);
@@ -652,7 +657,6 @@
                 final ActivityRecord ar = openApps.valueAt(i);
                 if (mAnimationHandler.isTarget(ar, true /* open */)) {
                     openApps.removeAt(i);
-                    mAnimationHandler.markStartingSurfaceMatch(null /* reparentTransaction */);
                 }
             }
             for (int i = closeApps.size() - 1; i >= 0; --i) {
@@ -670,6 +674,12 @@
         mAnimationHandler.markWindowHasDrawn(openActivity);
     }
 
+    boolean isStartingSurfaceShown(ActivityRecord openActivity) {
+        if (!Flags.migratePredictiveBackTransition()) {
+            return false;
+        }
+        return mAnimationHandler.isStartingSurfaceDrawn(openActivity);
+    }
     @VisibleForTesting
     class NavigationMonitor {
         // The window which triggering the back navigation.
@@ -767,8 +777,14 @@
      * open or close list.
      */
     void onTransactionReady(Transition transition, ArrayList<Transition.ChangeInfo> targets,
-            SurfaceControl.Transaction startTransaction) {
-        if (!isMonitoringTransition() || targets.isEmpty()) {
+            SurfaceControl.Transaction startTransaction,
+            SurfaceControl.Transaction finishTransaction) {
+        if (isMonitoringPrepareTransition(transition)) {
+            // Flag target matches and prepare to remove windowless surface.
+            mAnimationHandler.markStartingSurfaceMatch(startTransaction);
+            return;
+        }
+        if (!isMonitoringFinishTransition() || targets.isEmpty()) {
             return;
         }
         if (mAnimationHandler.hasTargetDetached()) {
@@ -801,42 +817,40 @@
         if (!matchAnimationTargets) {
             mNavigationMonitor.onTransitionReadyWhileNavigate(mTmpOpenApps, mTmpCloseApps);
         } else {
-            if (mWaitTransitionFinish != null) {
+            if (mAnimationHandler.mPrepareCloseTransition != null) {
                 Slog.e(TAG, "Gesture animation is applied on another transition?");
             }
-            mWaitTransitionFinish = transition;
-            // Flag target matches to defer remove the splash screen.
-            for (int i = mTmpOpenApps.size() - 1; i >= 0; --i) {
-                final WindowContainer wc = mTmpOpenApps.get(i);
-                if (mAnimationHandler.isTarget(wc, true /* open */)) {
-                    mAnimationHandler.markStartingSurfaceMatch(startTransaction);
-                    break;
-                }
+            mAnimationHandler.mPrepareCloseTransition = transition;
+            if (!Flags.migratePredictiveBackTransition()) {
+                // Because the target will reparent to transition root, so it cannot be controlled
+                // by animation leash. Hide the close target when transition starts.
+                startTransaction.hide(mAnimationHandler.mCloseAdaptor.mTarget.getSurfaceControl());
             }
+            // Flag target matches and prepare to remove windowless surface.
+            mAnimationHandler.markStartingSurfaceMatch(startTransaction);
             // release animation leash
             if (mAnimationHandler.mOpenAnimAdaptor.mCloseTransaction != null) {
-                startTransaction.merge(mAnimationHandler.mOpenAnimAdaptor.mCloseTransaction);
+                finishTransaction.merge(mAnimationHandler.mOpenAnimAdaptor.mCloseTransaction);
                 mAnimationHandler.mOpenAnimAdaptor.mCloseTransaction = null;
             }
-            // Because the target will reparent to transition root, so it cannot be controlled by
-            // animation leash. Hide the close target when transition starts.
-            startTransaction.hide(mAnimationHandler.mCloseAdaptor.mTarget.getSurfaceControl());
         }
         mTmpOpenApps.clear();
         mTmpCloseApps.clear();
     }
 
     boolean isMonitorTransitionTarget(WindowContainer wc) {
-        if (!isWaitBackTransition() || mWaitTransitionFinish == null) {
-            return false;
+        if ((isWaitBackTransition() && mAnimationHandler.mPrepareCloseTransition != null)
+                || (mAnimationHandler.mOpenAnimAdaptor != null
+                && mAnimationHandler.mOpenAnimAdaptor.mPreparedOpenTransition != null)) {
+            return mAnimationHandler.isTarget(wc, wc.isVisibleRequested() /* open */);
         }
-        return mAnimationHandler.isTarget(wc, wc.isVisibleRequested() /* open */);
+        return false;
     }
 
     boolean shouldPauseTouch(WindowContainer wc) {
         // Once the close transition is ready, it means the onBackInvoked callback has invoked, and
         // app is ready to trigger next transition, no matter what it will be.
-        return mAnimationHandler.mComposed && mWaitTransitionFinish == null
+        return mAnimationHandler.mComposed && mAnimationHandler.mPrepareCloseTransition == null
                 && mAnimationHandler.isTarget(wc, wc.isVisibleRequested() /* open */);
     }
 
@@ -847,7 +861,6 @@
     void clearBackAnimations(boolean cancel) {
         mAnimationHandler.clearBackAnimateTarget(cancel);
         mNavigationMonitor.stopMonitorTransition();
-        mWaitTransitionFinish = null;
     }
 
     /**
@@ -858,7 +871,13 @@
     */
     void onTransitionFinish(ArrayList<Transition.ChangeInfo> targets,
             @NonNull Transition finishedTransition) {
-        if (finishedTransition == mWaitTransitionFinish) {
+        if (isMonitoringPrepareTransition(finishedTransition)) {
+            if (mAnimationHandler.mPrepareCloseTransition == null) {
+                clearBackAnimations(true /* cancel */);
+            }
+            return;
+        }
+        if (finishedTransition == mAnimationHandler.mPrepareCloseTransition) {
             clearBackAnimations(false /* cancel */);
         }
         if (!mBackAnimationInProgress || mPendingAnimationBuilder == null) {
@@ -938,6 +957,7 @@
         // the opening target like starting window do.
         private boolean mStartingSurfaceTargetMatch;
         private ActivityRecord[] mOpenActivities;
+        Transition mPrepareCloseTransition;
 
         AnimationHandler(WindowManagerService wms) {
             mWindowManagerService = wms;
@@ -982,6 +1002,12 @@
                 @NonNull ActivityRecord[] openingActivities)  {
             if (isActivitySwitch(close, open)) {
                 mSwitchType = ACTIVITY_SWITCH;
+                if (Flags.migratePredictiveBackTransition()) {
+                    final Pair<WindowContainer, WindowContainer[]> replaced =
+                            promoteToTFIfNeeded(close, open);
+                    close = replaced.first;
+                    open = replaced.second;
+                }
             } else if (isTaskSwitch(close, open)) {
                 mSwitchType = TASK_SWITCH;
             } else if (isDialogClose(close)) {
@@ -1019,6 +1045,34 @@
             mOpenActivities = openingActivities;
         }
 
+        private Pair<WindowContainer, WindowContainer[]> promoteToTFIfNeeded(
+                WindowContainer close, WindowContainer[] open) {
+            WindowContainer replaceClose = close;
+            TaskFragment closeTF = close.asActivityRecord().getTaskFragment();
+            if (closeTF != null && !closeTF.isEmbedded()) {
+                closeTF = null;
+            }
+            final WindowContainer[] replaceOpen = new WindowContainer[open.length];
+            if (open.length >= 2) { // Promote to TaskFragment
+                for (int i = open.length - 1; i >= 0; --i) {
+                    replaceOpen[i] = open[i].asActivityRecord().getTaskFragment();
+                    replaceClose = closeTF != null ? closeTF : close;
+                }
+            } else {
+                TaskFragment openTF = open[0].asActivityRecord().getTaskFragment();
+                if (openTF != null && !openTF.isEmbedded()) {
+                    openTF = null;
+                }
+                if (closeTF != openTF) {
+                    replaceOpen[0] = openTF != null ? openTF : open[0];
+                    replaceClose = closeTF != null ? closeTF : close;
+                } else {
+                    replaceOpen[0] = open[0];
+                }
+            }
+            return new Pair<>(replaceClose, replaceOpen);
+        }
+
         private boolean composeAnimations(@NonNull WindowContainer close,
                 @NonNull WindowContainer[] open, @NonNull ActivityRecord[] openingActivities) {
             if (mComposed || mWaitTransition) {
@@ -1080,7 +1134,8 @@
         }
 
         void markWindowHasDrawn(ActivityRecord activity) {
-            if (!mComposed || mWaitTransition) {
+            if (!mComposed || mWaitTransition || mOpenAnimAdaptor.mPreparedOpenTransition == null
+                    || mOpenAnimAdaptor.mRequestedStartingSurfaceId == INVALID_TASK_ID) {
                 return;
             }
             boolean allWindowDrawn = true;
@@ -1096,6 +1151,17 @@
             }
         }
 
+        boolean isStartingSurfaceDrawn(ActivityRecord activity) {
+            // Check whether we create windowless surface to prepare open transition
+            if (!mComposed || mOpenAnimAdaptor.mPreparedOpenTransition == null) {
+                return false;
+            }
+            if (isTarget(activity, true /* open */)) {
+                return mOpenAnimAdaptor.mStartingSurface != null;
+            }
+            return false;
+        }
+
         private static boolean isAnimateTarget(@NonNull WindowContainer window,
                 @NonNull WindowContainer animationTarget, int switchType) {
             if (switchType == TASK_SWITCH) {
@@ -1109,7 +1175,9 @@
                         && window.hasChild(animationTarget));
             } else if (switchType == ACTIVITY_SWITCH) {
                 return window == animationTarget
-                        || (window.asTaskFragment() != null && window.hasChild(animationTarget));
+                        || (window.asTaskFragment() != null && window.hasChild(animationTarget))
+                        || (animationTarget.asTaskFragment() != null
+                        && animationTarget.hasChild(window));
             }
             return false;
         }
@@ -1122,8 +1190,14 @@
                         resetActivity.mDisplayContent
                                 .continueUpdateOrientationForDiffOrienLaunchingApp();
                     }
+                    final Transition finishTransition =
+                            resetActivity.mTransitionController.mFinishingTransition;
+                    final boolean inFinishTransition = finishTransition != null
+                            && (mPrepareCloseTransition == finishTransition
+                            || (mOpenAnimAdaptor != null
+                            && mOpenAnimAdaptor.mPreparedOpenTransition == finishTransition));
                     if (resetActivity.mLaunchTaskBehind) {
-                        restoreLaunchBehind(resetActivity, cancel);
+                        restoreLaunchBehind(resetActivity, cancel, inFinishTransition);
                     }
                 }
             }
@@ -1137,12 +1211,25 @@
             }
         }
 
-        void markStartingSurfaceMatch(SurfaceControl.Transaction reparentTransaction) {
+        void markStartingSurfaceMatch(SurfaceControl.Transaction startTransaction) {
             if (mStartingSurfaceTargetMatch) {
                 return;
             }
             mStartingSurfaceTargetMatch = true;
-            mOpenAnimAdaptor.reparentWindowlessSurfaceToTarget(reparentTransaction);
+
+            if (mOpenAnimAdaptor.mRequestedStartingSurfaceId == INVALID_TASK_ID) {
+                return;
+            }
+            final SurfaceControl startingSurface = mOpenAnimAdaptor.mStartingSurface;
+            if (startingSurface != null && startingSurface.isValid()) {
+                startTransaction.addTransactionCommittedListener(Runnable::run, () -> {
+                    synchronized (mWindowManagerService.mGlobalLock) {
+                        if (mOpenAnimAdaptor != null) {
+                            mOpenAnimAdaptor.cleanUpWindowlessSurface(true);
+                        }
+                    }
+                });
+            }
         }
 
         void clearBackAnimateTarget(boolean cancel) {
@@ -1150,6 +1237,7 @@
                 mComposed = false;
                 finishPresentAnimations(cancel);
             }
+            mPrepareCloseTransition = null;
             mWaitTransition = false;
             mStartingSurfaceTargetMatch = false;
             mSwitchType = UNKNOWN;
@@ -1239,6 +1327,9 @@
             // requested one during animating.
             private int mRequestedStartingSurfaceId = INVALID_TASK_ID;
             private SurfaceControl mStartingSurface;
+
+            private Transition mPreparedOpenTransition;
+
             BackWindowAnimationAdaptorWrapper(boolean isOpen, int switchType,
                     @NonNull WindowContainer... targets) {
                 mAdaptors = new BackWindowAnimationAdaptor[targets.length];
@@ -1267,6 +1358,8 @@
                     mCloseTransaction.apply();
                     mCloseTransaction = null;
                 }
+
+                mPreparedOpenTransition = null;
             }
 
             private RemoteAnimationTarget createWrapTarget() {
@@ -1279,8 +1372,7 @@
                     unionBounds.union(mAdaptors[i].mAnimationTarget.localBounds);
                 }
                 final WindowContainer wc = mAdaptors[0].mTarget;
-                final Task task = wc.asActivityRecord() != null
-                        ? wc.asActivityRecord().getTask() : wc.asTask();
+                final Task task = mAdaptors[0].getTopTask();
                 final RemoteAnimationTarget represent = mAdaptors[0].mAnimationTarget;
                 final SurfaceControl leashSurface = new SurfaceControl.Builder()
                         .setName("cross-animation-leash")
@@ -1293,7 +1385,7 @@
                 mCloseTransaction = new SurfaceControl.Transaction();
                 mCloseTransaction.reparent(leashSurface, null);
                 final SurfaceControl.Transaction pt = wc.getPendingTransaction();
-                pt.setLayer(leashSurface, wc.getParent().getLastLayer());
+                pt.setLayer(leashSurface, wc.getLastLayer());
                 for (int i = mAdaptors.length - 1; i >= 0; --i) {
                     BackWindowAnimationAdaptor adaptor = mAdaptors[i];
                     pt.reparent(adaptor.mAnimationTarget.leash, leashSurface);
@@ -1323,15 +1415,20 @@
                 }
                 final WindowContainer mainOpen = mAdaptors[0].mTarget;
                 final int switchType = mAdaptors[0].mSwitchType;
-                final Task openTask = switchType == TASK_SWITCH
-                        ? mainOpen.asTask() : switchType == ACTIVITY_SWITCH
-                        ? mainOpen.asActivityRecord().getTask() : null;
+                final Task openTask = mAdaptors[0].getTopTask();
                 if (openTask == null) {
                     return;
                 }
-                final ActivityRecord mainActivity = switchType == ACTIVITY_SWITCH
-                        ? mainOpen.asActivityRecord()
-                        : openTask.getTopNonFinishingActivity();
+                ActivityRecord mainActivity = null;
+                if (switchType == ACTIVITY_SWITCH) {
+                    mainActivity = mainOpen.asActivityRecord();
+                    if (mainActivity == null && mainOpen.asTaskFragment() != null) {
+                        mainActivity = mainOpen.asTaskFragment().getTopNonFinishingActivity();
+                    }
+                }
+                if (mainActivity == null) {
+                    mainActivity = openTask.getTopNonFinishingActivity();
+                }
                 if (mainActivity == null) {
                     return;
                 }
@@ -1363,6 +1460,8 @@
                                     synchronized (openTask.mWmService.mGlobalLock) {
                                         if (mRequestedStartingSurfaceId != INVALID_TASK_ID) {
                                             mStartingSurface = sc;
+                                            openTask.mWmService.mWindowPlacerLocked
+                                                    .requestTraversal();
                                         } else {
                                             sc.release();
                                         }
@@ -1371,28 +1470,6 @@
                             });
             }
 
-            // When back gesture has triggered and transition target matches navigation target,
-            // reparent the starting surface to the opening target as it's starting window.
-            void reparentWindowlessSurfaceToTarget(SurfaceControl.Transaction reparentTransaction) {
-                if (mRequestedStartingSurfaceId == INVALID_TASK_ID) {
-                    return;
-                }
-                // If open target matches, reparent to open activity or task
-                if (mStartingSurface != null && mStartingSurface.isValid()) {
-                    SurfaceControl.Transaction transaction = reparentTransaction != null
-                            ? reparentTransaction : mAdaptors[0].mTarget.getPendingTransaction();
-                    if (mAdaptors.length != 1) {
-                        // More than one opening window, reparent starting surface to leaf task.
-                        final WindowContainer wc = mAdaptors[0].mTarget;
-                        final Task task = wc.asActivityRecord() != null
-                                ? wc.asActivityRecord().getTask() : wc.asTask();
-                        transaction.reparent(mStartingSurface, task != null
-                                        ? task.getSurfaceControl()
-                                        : mAdaptors[0].mTarget.getSurfaceControl());
-                    }
-                }
-            }
-
             /**
              * Ask shell to clear the starting surface.
              * @param openTransitionMatch if true, shell will play the remove starting window
@@ -1430,6 +1507,22 @@
                 mSwitchType = switchType;
             }
 
+            Task getTopTask() {
+                final Task asTask = mTarget.asTask();
+                if (asTask != null) {
+                    return asTask;
+                }
+                final ActivityRecord ar = mTarget.asActivityRecord();
+                if (ar != null) {
+                    return ar.getTask();
+                }
+                final TaskFragment tf = mTarget.asTaskFragment();
+                if (tf != null) {
+                    return tf.getTask();
+                }
+                return null;
+            }
+
             @Override
             public boolean getShowWallpaper() {
                 return false;
@@ -1619,9 +1712,8 @@
                     needsLaunchBehind = snapshot == null;
                 }
                 if (needsLaunchBehind) {
-                    for (int i = visibleOpenActivities.length - 1; i >= 0; --i) {
-                        setLaunchBehind(visibleOpenActivities[i]);
-                    }
+                    openAnimationAdaptor.mPreparedOpenTransition =
+                            setLaunchBehind(visibleOpenActivities);
                 }
                 // Force update mLastSurfaceShowing for opening activity and its task.
                 if (mWindowManagerService.mRoot.mTransitionController.isShellTransitionsEnabled()) {
@@ -1677,13 +1769,22 @@
                                 // animation was canceled
                                 return;
                             }
-                            if (!triggerBack) {
-                                clearBackAnimateTarget(true /* cancel */);
+                            if (Flags.migratePredictiveBackTransition()) {
+                                if (mOpenAnimAdaptor == null
+                                        || mOpenAnimAdaptor.mPreparedOpenTransition == null) {
+                                    // no open nor close transition, this is window animation
+                                    if (!triggerBack) {
+                                        clearBackAnimateTarget(true /* cancel */);
+                                    }
+                                }
                             } else {
-                                mWaitTransition = true;
+                                if (!triggerBack) {
+                                    clearBackAnimateTarget(true /* cancel */);
+                                } else {
+                                    mWaitTransition = true;
+                                }
                             }
                         }
-                        // TODO Add timeout monitor if transition didn't happen
                     }
                 };
             }
@@ -1740,28 +1841,75 @@
         return openActivities;
     }
 
-    private static void setLaunchBehind(@NonNull ActivityRecord activity) {
-        if (!activity.isVisibleRequested()) {
-            // The transition could commit the visibility and in the finishing state, that could
-            // skip commitVisibility call in setVisibility cause the activity won't visible here.
-            // Call it again to make sure the activity could be visible while handling the pending
-            // animation.
-            // Do not performLayout during prepare animation, because it could cause focus window
-            // change. Let that happen after the BackNavigationInfo has returned to shell.
-            activity.commitVisibility(true, false /* performLayout */);
+    private static Transition setLaunchBehind(@NonNull ActivityRecord[] activities) {
+        final boolean migrateBackTransition = Flags.migratePredictiveBackTransition();
+        final ArrayList<ActivityRecord> affects = new ArrayList<>();
+        for (int i = activities.length - 1; i >= 0; --i) {
+            final ActivityRecord activity = activities[i];
+            if (activity.mLaunchTaskBehind || activity.isVisibleRequested()) {
+                continue;
+            }
+            affects.add(activity);
+        }
+        if (affects.isEmpty()) {
+            return null;
+        }
+
+        final TransitionController tc = activities[0].mTransitionController;
+        final Transition prepareOpen = migrateBackTransition && !tc.isCollecting()
+                ? tc.createTransition(TRANSIT_PREPARE_BACK_NAVIGATION) : null;
+
+        for (int i = affects.size() - 1; i >= 0; --i) {
+            final ActivityRecord activity = affects.get(i);
+            if (!migrateBackTransition && !activity.isVisibleRequested()) {
+                // The transition could commit the visibility and in the finishing state, that could
+                // skip commitVisibility call in setVisibility cause the activity won't visible
+                // here.
+                // Call it again to make sure the activity could be visible while handling the
+                // pending animation.
+                // Do not performLayout during prepare animation, because it could cause focus
+                // window change. Let that happen after the BackNavigationInfo has returned to
+                // shell.
+                activity.commitVisibility(true, false /* performLayout */);
+            }
             activity.mTransitionController.mSnapshotController
                     .mActivitySnapshotController.addOnBackPressedActivity(activity);
-        }
-        activity.mLaunchTaskBehind = true;
+            activity.mLaunchTaskBehind = true;
 
-        ProtoLog.d(WM_DEBUG_BACK_PREVIEW,
-                "Setting Activity.mLauncherTaskBehind to true. Activity=%s", activity);
-        activity.mTaskSupervisor.mStoppingActivities.remove(activity);
-        activity.getDisplayContent().ensureActivitiesVisible(null /* starting */,
-                true /* notifyClients */);
+            ProtoLog.d(WM_DEBUG_BACK_PREVIEW,
+                    "Setting Activity.mLauncherTaskBehind to true. Activity=%s", activity);
+            activity.mTaskSupervisor.mStoppingActivities.remove(activity);
+
+            if (!migrateBackTransition) {
+                activity.getDisplayContent().ensureActivitiesVisible(null /* starting */,
+                        true /* notifyClients */);
+            } else if (activity.shouldBeVisible()) {
+                activity.ensureActivityConfiguration(true /* ignoreVisibility */);
+                activity.makeVisibleIfNeeded(null /* starting */, true /* notifyToClient */);
+            }
+        }
+        boolean needTransition = false;
+        final DisplayContent dc = affects.get(0).getDisplayContent();
+        for (int i = affects.size() - 1; i >= 0; --i) {
+            final ActivityRecord activity = affects.get(i);
+            needTransition |= tc.isCollecting(activity);
+        }
+        if (prepareOpen != null) {
+            if (needTransition) {
+                tc.requestStartTransition(prepareOpen,
+                        null /*startTask */, null /* remoteTransition */,
+                        null /* displayChange */);
+                tc.setReady(dc);
+                return prepareOpen;
+            } else {
+                prepareOpen.abort();
+            }
+        }
+        return null;
     }
 
-    private static void restoreLaunchBehind(@NonNull ActivityRecord activity, boolean cancel) {
+    private static void restoreLaunchBehind(@NonNull ActivityRecord activity, boolean cancel,
+            boolean finishTransition) {
         if (!activity.isAttached()) {
             // The activity was detached from hierarchy.
             return;
@@ -1771,9 +1919,15 @@
                 "Setting Activity.mLauncherTaskBehind to false. Activity=%s",
                 activity);
         if (cancel) {
-            // Restore the launch-behind state
-            // TODO b/347168362 Change status directly during collecting for a transition.
-            activity.mTaskSupervisor.scheduleLaunchTaskBehindComplete(activity.token);
+            final boolean migrateBackTransition = Flags.migratePredictiveBackTransition();
+            if (migrateBackTransition && finishTransition) {
+                activity.commitVisibility(false /* visible */, false /* performLayout */,
+                        true /* fromTransition */);
+            } else {
+                // Restore the launch-behind state
+                // TODO b/347168362 Change status directly during collecting for a transition.
+                activity.mTaskSupervisor.scheduleLaunchTaskBehindComplete(activity.token);
+            }
             // Ignore all change
             activity.mTransitionController.mSnapshotController
                     .mActivitySnapshotController.clearOnBackPressedActivities();
@@ -1835,12 +1989,7 @@
                         && ah.mSwitchType != AnimationHandler.ACTIVITY_SWITCH)) {
                     return;
                 }
-                for (int i = mAnimationHandler.mOpenActivities.length - 1; i >= 0; --i) {
-                    final ActivityRecord preDrawActivity = mAnimationHandler.mOpenActivities[i];
-                    if (!preDrawActivity.mLaunchTaskBehind) {
-                        setLaunchBehind(preDrawActivity);
-                    }
-                }
+                setLaunchBehind(mAnimationHandler.mOpenActivities);
             }
         }
     }
@@ -1853,10 +2002,15 @@
             snapshot = task.mRootWindowContainer.mWindowManager.mTaskSnapshotController.getSnapshot(
                     task.mTaskId, task.mUserId, false /* restoreFromDisk */,
                     false /* isLowResolution */);
-        } else if (w.asActivityRecord() != null) {
-            final ActivityRecord ar = w.asActivityRecord();
-            snapshot = ar.mWmService.mSnapshotController.mActivitySnapshotController
-                    .getSnapshot(visibleOpenActivities);
+        } else {
+            ActivityRecord ar = w.asActivityRecord();
+            if (ar == null && w.asTaskFragment() != null) {
+                ar = w.asTaskFragment().getTopNonFinishingActivity();
+            }
+            if (ar != null) {
+                snapshot = ar.mWmService.mSnapshotController.mActivitySnapshotController
+                        .getSnapshot(visibleOpenActivities);
+            }
         }
 
         return isSnapshotCompatible(snapshot, visibleOpenActivities) ? snapshot : null;
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 996b439..5698750 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1764,7 +1764,7 @@
 
         // Check whether the participants were animated from back navigation.
         mController.mAtm.mBackNavigationController.onTransactionReady(this, mTargets,
-                transaction);
+                transaction, mFinishTransaction);
         final TransitionInfo info = calculateTransitionInfo(mType, mFlags, mTargets, transaction);
         info.setDebugId(mSyncId);
         mController.assignTrack(this, info);
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
index a159ce3..44837d7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
@@ -58,6 +58,7 @@
 import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
+import android.platform.test.annotations.RequiresFlagsDisabled;
 import android.util.ArraySet;
 import android.view.WindowManager;
 import android.window.BackAnimationAdapter;
@@ -72,6 +73,7 @@
 import android.window.WindowOnBackInvokedDispatcher;
 
 import com.android.server.LocalServices;
+import com.android.window.flags.Flags;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -612,6 +614,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_MIGRATE_PREDICTIVE_BACK_TRANSITION)
     public void testTransitionHappensCancelNavigation() {
         // Create a floating task and a fullscreen task, then navigating on fullscreen task.
         // The navigation should not been cancelled when transition happens on floating task, and