Merge "Import SettingsLib/Metadata library" into main
diff --git a/core/java/android/window/TaskFragmentOrganizer.java b/core/java/android/window/TaskFragmentOrganizer.java
index 4cc0d8a..c3168001 100644
--- a/core/java/android/window/TaskFragmentOrganizer.java
+++ b/core/java/android/window/TaskFragmentOrganizer.java
@@ -69,6 +69,23 @@
     public static final String KEY_ERROR_CALLBACK_OP_TYPE = "operation_type";
 
     /**
+     * Key to bundle {@link TaskFragmentInfo}s from the system in
+     * {@link #registerOrganizer(boolean, Bundle)}
+     *
+     * @hide
+     */
+    public static final String KEY_RESTORE_TASK_FRAGMENTS_INFO = "key_restore_task_fragments_info";
+
+    /**
+     * Key to bundle {@link TaskFragmentParentInfo} from the system in
+     * {@link #registerOrganizer(boolean, Bundle)}
+     *
+     * @hide
+     */
+    public static final String KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO =
+            "key_restore_task_fragment_parent_info";
+
+    /**
      * No change set.
      */
     @WindowManager.TransitionType
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
index 4ce2942..bfccb29 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
@@ -16,16 +16,26 @@
 
 package androidx.window.extensions.embedding;
 
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENTS_INFO;
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO;
+
 import android.os.Build;
 import android.os.Bundle;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.MessageQueue;
+import android.util.ArrayMap;
 import android.util.Log;
+import android.util.SparseArray;
+import android.window.TaskFragmentInfo;
+import android.window.TaskFragmentParentInfo;
+import android.window.WindowContainerTransaction;
 
 import androidx.annotation.NonNull;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 
 /**
  * Helper class to back up and restore the TaskFragmentOrganizer state, in order to resume
@@ -40,11 +50,21 @@
     @NonNull
     private final SplitController mController;
     @NonNull
+    private final SplitPresenter mPresenter;
+    @NonNull
     private final BackupIdler mBackupIdler = new BackupIdler();
     private boolean mBackupIdlerScheduled;
 
-    BackupHelper(@NonNull SplitController splitController, @NonNull Bundle savedState) {
+    private final List<ParcelableTaskContainerData> mParcelableTaskContainerDataList =
+            new ArrayList<>();
+    private final ArrayMap<IBinder, TaskFragmentInfo> mTaskFragmentInfos = new ArrayMap<>();
+    private final SparseArray<TaskFragmentParentInfo> mTaskFragmentParentInfos =
+            new SparseArray<>();
+
+    BackupHelper(@NonNull SplitController splitController, @NonNull SplitPresenter splitPresenter,
+            @NonNull Bundle savedState) {
         mController = splitController;
+        mPresenter = splitPresenter;
 
         if (!savedState.isEmpty()) {
             restoreState(savedState);
@@ -67,13 +87,13 @@
         public boolean queueIdle() {
             synchronized (mController.mLock) {
                 mBackupIdlerScheduled = false;
-                startBackup();
+                saveState();
             }
             return false;
         }
     }
 
-    private void startBackup() {
+    private void saveState() {
         final List<TaskContainer> taskContainers = mController.getTaskContainers();
         if (taskContainers.isEmpty()) {
             Log.w(TAG, "No task-container to back up");
@@ -97,13 +117,92 @@
             return;
         }
 
-        final List<ParcelableTaskContainerData> parcelableTaskContainerDataList =
-                savedState.getParcelableArrayList(KEY_TASK_CONTAINERS,
-                        ParcelableTaskContainerData.class);
-        for (ParcelableTaskContainerData data : parcelableTaskContainerDataList) {
-            final TaskContainer taskContainer = new TaskContainer(data, mController);
-            if (DEBUG) Log.d(TAG, "Restoring task " + taskContainer.getTaskId());
-            // TODO(b/289875940): implement the TaskContainer restoration.
+        if (DEBUG) Log.d(TAG, "Start restoring saved-state");
+        mParcelableTaskContainerDataList.addAll(savedState.getParcelableArrayList(
+                KEY_TASK_CONTAINERS, ParcelableTaskContainerData.class));
+        if (DEBUG) Log.d(TAG, "Retrieved tasks : " + mParcelableTaskContainerDataList.size());
+        if (mParcelableTaskContainerDataList.isEmpty()) {
+            return;
+        }
+
+        final List<TaskFragmentInfo> infos = savedState.getParcelableArrayList(
+                KEY_RESTORE_TASK_FRAGMENTS_INFO, TaskFragmentInfo.class);
+        for (TaskFragmentInfo info : infos) {
+            if (DEBUG) Log.d(TAG, "Retrieved: " + info);
+            mTaskFragmentInfos.put(info.getFragmentToken(), info);
+            mPresenter.updateTaskFragmentInfo(info);
+        }
+
+        final List<TaskFragmentParentInfo> parentInfos = savedState.getParcelableArrayList(
+                KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO,
+                TaskFragmentParentInfo.class);
+        for (TaskFragmentParentInfo info : parentInfos) {
+            if (DEBUG) Log.d(TAG, "Retrieved: " + info);
+            mTaskFragmentParentInfos.put(info.getTaskId(), info);
         }
     }
-}
+
+    boolean hasPendingStateToRestore() {
+        return !mParcelableTaskContainerDataList.isEmpty();
+    }
+
+    /**
+     * Returns {@code true} if any of the {@link TaskContainer} is restored.
+     * Otherwise, returns {@code false}.
+     */
+    boolean rebuildTaskContainers(@NonNull WindowContainerTransaction wct,
+            @NonNull Set<EmbeddingRule> rules) {
+        if (mParcelableTaskContainerDataList.isEmpty()) {
+            return false;
+        }
+
+        if (DEBUG) Log.d(TAG, "Rebuilding TaskContainers.");
+        final ArrayMap<String, EmbeddingRule> embeddingRuleMap = new ArrayMap<>();
+        for (EmbeddingRule rule : rules) {
+            embeddingRuleMap.put(rule.getTag(), rule);
+        }
+
+        boolean restoredAny = false;
+        for (int i = mParcelableTaskContainerDataList.size() - 1; i >= 0; i--) {
+            final ParcelableTaskContainerData parcelableTaskContainerData =
+                    mParcelableTaskContainerDataList.get(i);
+            final List<String> tags = parcelableTaskContainerData.getSplitRuleTags();
+            if (!embeddingRuleMap.containsAll(tags)) {
+                // has unknown tag, unable to restore.
+                if (DEBUG) {
+                    Log.d(TAG, "Rebuilding TaskContainer abort! Unknown Tag. Task#"
+                            + parcelableTaskContainerData.mTaskId);
+                }
+                continue;
+            }
+
+            mParcelableTaskContainerDataList.remove(parcelableTaskContainerData);
+            final TaskContainer taskContainer = new TaskContainer(parcelableTaskContainerData,
+                    mController, mTaskFragmentInfos);
+            if (DEBUG) Log.d(TAG, "Created TaskContainer " + taskContainer);
+            mController.addTaskContainer(taskContainer.getTaskId(), taskContainer);
+
+            for (ParcelableSplitContainerData splitData :
+                    parcelableTaskContainerData.getParcelableSplitContainerDataList()) {
+                final SplitRule rule = (SplitRule) embeddingRuleMap.get(splitData.mSplitRuleTag);
+                assert rule != null;
+                if (mController.getContainer(splitData.getPrimaryContainerToken()) != null
+                        && mController.getContainer(splitData.getSecondaryContainerToken())
+                        != null) {
+                    taskContainer.addSplitContainer(
+                            new SplitContainer(splitData, mController, rule));
+                }
+            }
+
+            mController.onTaskFragmentParentRestored(wct, taskContainer.getTaskId(),
+                    mTaskFragmentParentInfos.get(taskContainer.getTaskId()));
+            restoredAny = true;
+        }
+
+        if (mParcelableTaskContainerDataList.isEmpty()) {
+            mTaskFragmentParentInfos.clear();
+            mTaskFragmentInfos.clear();
+        }
+        return restoredAny;
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
index 817cfce..cb280c5 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
@@ -89,13 +89,13 @@
     };
 
     @NonNull
-    private IBinder getPrimaryContainerToken() {
+    IBinder getPrimaryContainerToken() {
         return mSplitContainer != null ? mSplitContainer.getPrimaryContainer().getToken()
                 : mPrimaryContainerToken;
     }
 
     @NonNull
-    private IBinder getSecondaryContainerToken() {
+    IBinder getSecondaryContainerToken() {
         return mSplitContainer != null ? mSplitContainer.getSecondaryContainer().getToken()
                 : mSecondaryContainerToken;
     }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
index 7377d00..97aa699 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
@@ -108,6 +108,15 @@
                 : mParcelableSplitContainerDataList;
     }
 
+    @NonNull
+    List<String> getSplitRuleTags() {
+        final List<String> tags = new ArrayList<>();
+        for (ParcelableSplitContainerData data : getParcelableSplitContainerDataList()) {
+            tags.add(data.mSplitRuleTag);
+        }
+        return tags;
+    }
+
     @Override
     public int describeContents() {
         return 0;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
index 6d436ec..faf73c2 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
@@ -86,6 +86,25 @@
         }
     }
 
+    /** This is only used when restoring it from a {@link ParcelableSplitContainerData}. */
+    SplitContainer(@NonNull ParcelableSplitContainerData parcelableData,
+            @NonNull SplitController splitController, @NonNull SplitRule splitRule) {
+        mParcelableData = parcelableData;
+        mPrimaryContainer = splitController.getContainer(parcelableData.getPrimaryContainerToken());
+        mSecondaryContainer = splitController.getContainer(
+                parcelableData.getSecondaryContainerToken());
+        mSplitRule = splitRule;
+        mDefaultSplitAttributes = splitRule.getDefaultSplitAttributes();
+        mCurrentSplitAttributes = mDefaultSplitAttributes;
+
+        if (shouldFinishPrimaryWithSecondary(splitRule)) {
+            mSecondaryContainer.addContainerToFinishOnExit(mPrimaryContainer);
+        }
+        if (shouldFinishSecondaryWithPrimary(splitRule)) {
+            mPrimaryContainer.addContainerToFinishOnExit(mSecondaryContainer);
+        }
+    }
+
     void setPrimaryContainer(@NonNull TaskFragmentContainer primaryContainer) {
         if (!mParcelableData.mIsPrimaryContainerMutable) {
             throw new IllegalStateException("Cannot update primary TaskFragmentContainer");
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index f2f2b7ea..db4bb0e 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -279,6 +279,26 @@
             Log.i(TAG, "Setting embedding rules. Size: " + rules.size());
             mSplitRules.clear();
             mSplitRules.addAll(rules);
+
+            if (!Flags.aeBackStackRestore() || !mPresenter.isRebuildTaskContainersNeeded()) {
+                return;
+            }
+
+            try {
+                final TransactionRecord transactionRecord =
+                        mTransactionManager.startNewTransaction();
+                final WindowContainerTransaction wct = transactionRecord.getTransaction();
+                if (mPresenter.rebuildTaskContainers(wct, rules)) {
+                    transactionRecord.apply(false /* shouldApplyIndependently */);
+                    updateCallbackIfNecessary();
+                } else {
+                    transactionRecord.abort();
+                }
+            } catch (IllegalStateException ex) {
+                Log.e(TAG, "Having an existing transaction while running restoration with"
+                        + "new rules!! It is likely too late to perform the restoration "
+                        + "already!?", ex);
+            }
         }
     }
 
@@ -903,6 +923,12 @@
     }
 
     @GuardedBy("mLock")
+    void onTaskFragmentParentRestored(@NonNull WindowContainerTransaction wct, int taskId,
+            @NonNull TaskFragmentParentInfo parentInfo) {
+        onTaskFragmentParentInfoChanged(wct, taskId, parentInfo);
+    }
+
+    @GuardedBy("mLock")
     void updateContainersInTaskIfVisible(@NonNull WindowContainerTransaction wct, int taskId) {
         final TaskContainer taskContainer = getTaskContainer(taskId);
         if (taskContainer == null) {
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index abc7b29..0c0ded9 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -24,6 +24,7 @@
 import static androidx.window.extensions.embedding.WindowAttributes.DIM_AREA_ON_TASK;
 
 import android.annotation.AnimRes;
+import android.annotation.NonNull;
 import android.app.Activity;
 import android.app.ActivityThread;
 import android.app.WindowConfiguration;
@@ -47,7 +48,6 @@
 import android.window.WindowContainerTransaction;
 
 import androidx.annotation.IntDef;
-import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.window.extensions.core.util.function.Function;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType;
@@ -67,6 +67,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.Executor;
 
 /**
@@ -174,7 +175,7 @@
         } else {
             registerOrganizer();
         }
-        mBackupHelper = new BackupHelper(controller, outSavedState);
+        mBackupHelper = new BackupHelper(controller, this, outSavedState);
         if (!SplitController.ENABLE_SHELL_TRANSITIONS) {
             // TODO(b/207070762): cleanup with legacy app transition
             // Animation will be handled by WM Shell when Shell transition is enabled.
@@ -186,6 +187,15 @@
         mBackupHelper.scheduleBackup();
     }
 
+    boolean isRebuildTaskContainersNeeded() {
+        return mBackupHelper.hasPendingStateToRestore();
+    }
+
+    boolean rebuildTaskContainers(@NonNull WindowContainerTransaction wct,
+            @NonNull Set<EmbeddingRule> rules) {
+        return mBackupHelper.rebuildTaskContainers(wct, rules);
+    }
+
     /**
      * Deletes the specified container and all other associated and dependent containers in the same
      * transaction.
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
index 608a3be..74cce68 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
@@ -31,6 +31,7 @@
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.IBinder;
+import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.DisplayMetrics;
 import android.util.Log;
@@ -147,14 +148,23 @@
 
     /** This is only used when restoring it from a {@link ParcelableTaskContainerData}. */
     TaskContainer(@NonNull ParcelableTaskContainerData data,
-            @NonNull SplitController splitController) {
+            @NonNull SplitController splitController,
+            @NonNull ArrayMap<IBinder, TaskFragmentInfo> taskFragmentInfoMap) {
         mParcelableTaskContainerData = new ParcelableTaskContainerData(data, this);
+        mInfo = new TaskFragmentParentInfo(new Configuration(), 0 /* displayId */, -1 /* taskId */,
+                false /* visible */, false /* hasDirectActivity */, null /* decorSurface */);
         mSplitController = splitController;
         for (ParcelableTaskFragmentContainerData tfData :
                 data.getParcelableTaskFragmentContainerDataList()) {
-            final TaskFragmentContainer container =
-                    new TaskFragmentContainer(tfData, splitController, this);
-            mContainers.add(container);
+            final TaskFragmentInfo info = taskFragmentInfoMap.get(tfData.mToken);
+            if (info != null && !info.isEmpty()) {
+                final TaskFragmentContainer container =
+                        new TaskFragmentContainer(tfData, splitController, this);
+                container.setInfo(new WindowContainerTransaction(), info);
+                mContainers.add(container);
+            } else {
+                Log.d(TAG, "Drop " + tfData + " while restoring Task " + data.mTaskId);
+            }
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java
index 3d79a1c..c385f9a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java
@@ -74,8 +74,12 @@
             final boolean isBackGesture = change.hasFlags(FLAG_BACK_GESTURE_ANIMATED);
             if (taskInfo.getActivityType() == ACTIVITY_TYPE_HOME) {
                 if (Flags.migratePredictiveBackTransition()) {
-                    if (!isBackGesture && TransitionUtil.isOpenOrCloseMode(mode)) {
-                        notifyHomeVisibilityChanged(TransitionUtil.isOpeningType(mode));
+                    final boolean gestureToHomeTransition = isBackGesture
+                            && TransitionUtil.isClosingType(info.getType());
+                    if (gestureToHomeTransition
+                            || (!isBackGesture && TransitionUtil.isOpenOrCloseMode(mode))) {
+                        notifyHomeVisibilityChanged(gestureToHomeTransition
+                                || TransitionUtil.isOpeningType(mode));
                     }
                 } else {
                     if (TransitionUtil.isOpenOrCloseMode(mode) || isBackGesture) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java
index f51a960..8f49de0 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java
@@ -20,6 +20,7 @@
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED;
 import static android.view.WindowManager.TRANSIT_CHANGE;
 import static android.view.WindowManager.TRANSIT_OPEN;
+import static android.view.WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.window.TransitionInfo.FLAG_BACK_GESTURE_ANIMATED;
 
@@ -39,6 +40,7 @@
 import android.os.Looper;
 import android.os.RemoteException;
 import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.view.SurfaceControl;
@@ -213,6 +215,35 @@
         verify(mListener, times(1)).onHomeVisibilityChanged(true);
     }
 
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_MIGRATE_PREDICTIVE_BACK_TRANSITION)
+    public void testHomeActivityWithBackGestureNotifiesHomeIsVisibleAfterClose()
+            throws RemoteException {
+        TransitionInfo info = mock(TransitionInfo.class);
+        TransitionInfo.Change change = mock(TransitionInfo.Change.class);
+        ActivityManager.RunningTaskInfo taskInfo = mock(ActivityManager.RunningTaskInfo.class);
+        when(change.getTaskInfo()).thenReturn(taskInfo);
+        when(info.getChanges()).thenReturn(new ArrayList<>(List.of(change)));
+        when(info.getType()).thenReturn(TRANSIT_PREPARE_BACK_NAVIGATION);
+
+        when(change.hasFlags(FLAG_BACK_GESTURE_ANIMATED)).thenReturn(true);
+        setupTransitionInfo(taskInfo, change, ACTIVITY_TYPE_HOME, TRANSIT_OPEN, true);
+
+        mHomeTransitionObserver.onTransitionReady(mock(IBinder.class),
+                info,
+                mock(SurfaceControl.Transaction.class),
+                mock(SurfaceControl.Transaction.class));
+        verify(mListener, times(0)).onHomeVisibilityChanged(anyBoolean());
+
+        when(info.getType()).thenReturn(TRANSIT_TO_BACK);
+        setupTransitionInfo(taskInfo, change, ACTIVITY_TYPE_HOME, TRANSIT_CHANGE, true);
+        mHomeTransitionObserver.onTransitionReady(mock(IBinder.class),
+                info,
+                mock(SurfaceControl.Transaction.class),
+                mock(SurfaceControl.Transaction.class));
+        verify(mListener, times(1)).onHomeVisibilityChanged(true);
+    }
+
     /**
      * Helper class to initialize variables for the rest.
      */
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/MirrorWindowControlTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/MirrorWindowControlTest.java
index 8f9b7c8..12c866f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/MirrorWindowControlTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/MirrorWindowControlTest.java
@@ -30,11 +30,11 @@
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
-import android.view.WindowManager;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.systemui.SysuiTestCase;
 
 import org.junit.Before;
@@ -48,7 +48,7 @@
 @RunWith(AndroidJUnit4.class)
 public class MirrorWindowControlTest extends SysuiTestCase {
 
-    @Mock WindowManager mWindowManager;
+    @Mock ViewCaptureAwareWindowManager mWindowManager;
     View mView;
     int mViewWidth;
     int mViewHeight;
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/MirrorWindowControl.java b/packages/SystemUI/src/com/android/systemui/accessibility/MirrorWindowControl.java
index 443441f..eb4de68 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/MirrorWindowControl.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/MirrorWindowControl.java
@@ -18,6 +18,9 @@
 
 import static android.view.WindowManager.LayoutParams;
 
+import static com.android.app.viewcapture.ViewCaptureFactory.getViewCaptureAwareWindowManagerInstance;
+import static com.android.systemui.Flags.enableViewCaptureTracing;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
@@ -29,8 +32,8 @@
 import android.view.Gravity;
 import android.view.LayoutInflater;
 import android.view.View;
-import android.view.WindowManager;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.systemui.res.R;
 
 /**
@@ -70,11 +73,12 @@
      * @see #setDefaultPosition(LayoutParams)
      */
     private final Point mControlPosition = new Point();
-    private final WindowManager mWindowManager;
+    private final ViewCaptureAwareWindowManager mWindowManager;
 
     MirrorWindowControl(Context context) {
         mContext = context;
-        mWindowManager = (WindowManager) mContext.getSystemService(Context.WINDOW_SERVICE);
+        mWindowManager = getViewCaptureAwareWindowManagerInstance(mContext,
+                enableViewCaptureTracing());
     }
 
     public void setWindowDelegate(@Nullable MirrorWindowDelegate windowDelegate) {
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 1fcd7f1..ed9dcfa 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -535,7 +535,7 @@
     volatile boolean mRequestedOrSleepingDefaultDisplay;
 
     /**
-     * This is used to check whether to invoke {@link #updateScreenOffSleepToken} when screen is
+     * This is used to check whether to acquire screen-off sleep token when screen is
      * turned off. E.g. if it is false when screen is turned off and the display is swapping, it
      * is expected that the screen will be on in a short time. Then it is unnecessary to acquire
      * screen-off-sleep-token, so it can avoid intermediate visibility or lifecycle changes.
@@ -610,7 +610,6 @@
     private boolean mPendingKeyguardOccluded;
     private boolean mKeyguardOccludedChanged;
 
-    private ActivityTaskManagerInternal.SleepTokenAcquirer mScreenOffSleepTokenAcquirer;
     Intent mHomeIntent;
     Intent mCarDockIntent;
     Intent mDeskDockIntent;
@@ -2220,9 +2219,6 @@
         mLockPatternUtils = new LockPatternUtils(mContext);
         mLogger = new MetricsLogger();
 
-        mScreenOffSleepTokenAcquirer = mActivityTaskManagerInternal
-                .createSleepTokenAcquirer("ScreenOff");
-
         Resources res = mContext.getResources();
         mWakeOnDpadKeyPress =
                 res.getBoolean(com.android.internal.R.bool.config_wakeOnDpadKeyPress);
@@ -5521,13 +5517,6 @@
         mRequestedOrSleepingDefaultDisplay = true;
         mIsGoingToSleepDefaultDisplay = true;
 
-        // In case startedGoingToSleep is called after screenTurnedOff (the source caller is in
-        // order but the methods run on different threads) and updateScreenOffSleepToken was
-        // skipped. Then acquire sleep token if screen was off.
-        if (!mDefaultDisplayPolicy.isScreenOnFully() && !mDefaultDisplayPolicy.isScreenOnEarly()) {
-            updateScreenOffSleepToken(true /* acquire */);
-        }
-
         if (mKeyguardDelegate != null) {
             mKeyguardDelegate.onStartedGoingToSleep(pmSleepReason);
         }
@@ -5688,11 +5677,9 @@
         if (DEBUG_WAKEUP) Slog.i(TAG, "Display" + displayId + " turned off...");
 
         if (displayId == DEFAULT_DISPLAY) {
-            if (!isSwappingDisplay || mIsGoingToSleepDefaultDisplay) {
-                updateScreenOffSleepToken(true /* acquire */);
-            }
+            final boolean acquireSleepToken = !isSwappingDisplay || mIsGoingToSleepDefaultDisplay;
             mRequestedOrSleepingDefaultDisplay = false;
-            mDefaultDisplayPolicy.screenTurnedOff();
+            mDefaultDisplayPolicy.screenTurnedOff(acquireSleepToken);
             synchronized (mLock) {
                 if (mKeyguardDelegate != null) {
                     mKeyguardDelegate.onScreenTurnedOff();
@@ -5748,7 +5735,6 @@
         if (displayId == DEFAULT_DISPLAY) {
             Trace.asyncTraceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "screenTurningOn",
                     0 /* cookie */);
-            updateScreenOffSleepToken(false /* acquire */);
             mDefaultDisplayPolicy.screenTurningOn(screenOnListener);
             mBootAnimationDismissable = false;
 
@@ -6255,15 +6241,6 @@
         }
     }
 
-    // TODO (multidisplay): Support multiple displays in WindowManagerPolicy.
-    private void updateScreenOffSleepToken(boolean acquire) {
-        if (acquire) {
-            mScreenOffSleepTokenAcquirer.acquire(DEFAULT_DISPLAY);
-        } else {
-            mScreenOffSleepTokenAcquirer.release(DEFAULT_DISPLAY);
-        }
-    }
-
     /** {@inheritDoc} */
     @Override
     public void enableScreenAfterBoot() {
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
index b8ce02ed..3d6b64b 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
@@ -130,35 +130,6 @@
     }
 
     /**
-     * Sleep tokens cause the activity manager to put the top activity to sleep.
-     * They are used by components such as dreams that may hide and block interaction
-     * with underlying activities.
-     * The Acquirer provides an interface that encapsulates the underlying work, so the user does
-     * not need to handle the token by him/herself.
-     */
-    public interface SleepTokenAcquirer {
-
-        /**
-         * Acquires a sleep token.
-         * @param displayId The display to apply to.
-         */
-        void acquire(int displayId);
-
-        /**
-         * Releases the sleep token.
-         * @param displayId The display to apply to.
-         */
-        void release(int displayId);
-    }
-
-    /**
-     * Creates a sleep token acquirer for the specified display with the specified tag.
-     *
-     * @param tag A string identifying the purpose (eg. "Dream").
-     */
-    public abstract SleepTokenAcquirer createSleepTokenAcquirer(@NonNull String tag);
-
-    /**
      * Returns home activity for the specified user.
      *
      * @param userId ID of the user or {@link android.os.UserHandle#USER_ALL}
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index e25d940d..49ca698 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -4356,6 +4356,7 @@
             mTaskOrganizerController.dump(pw, "  ");
             mVisibleActivityProcessTracker.dump(pw, "  ");
             mActiveUids.dump(pw, "  ");
+            pw.println("  SleepTokens=" + mRootWindowContainer.mSleepTokens);
             if (mDemoteTopAppReasons != 0) {
                 pw.println("  mDemoteTopAppReasons=" + mDemoteTopAppReasons);
             }
@@ -5071,17 +5072,16 @@
         EventLogTags.writeWmSetResumedActivity(r.mUserId, r.shortComponentName, reason);
     }
 
-    final class SleepTokenAcquirerImpl implements ActivityTaskManagerInternal.SleepTokenAcquirer {
+    final class SleepTokenAcquirer {
         private final String mTag;
         private final SparseArray<RootWindowContainer.SleepToken> mSleepTokens =
                 new SparseArray<>();
 
-        SleepTokenAcquirerImpl(@NonNull String tag) {
+        SleepTokenAcquirer(@NonNull String tag) {
             mTag = tag;
         }
 
-        @Override
-        public void acquire(int displayId) {
+        void acquire(int displayId) {
             synchronized (mGlobalLock) {
                 if (!mSleepTokens.contains(displayId)) {
                     mSleepTokens.append(displayId,
@@ -5091,8 +5091,7 @@
             }
         }
 
-        @Override
-        public void release(int displayId) {
+        void release(int displayId) {
             synchronized (mGlobalLock) {
                 final RootWindowContainer.SleepToken token = mSleepTokens.get(displayId);
                 if (token != null) {
@@ -5955,11 +5954,6 @@
     }
 
     final class LocalService extends ActivityTaskManagerInternal {
-        @Override
-        public SleepTokenAcquirer createSleepTokenAcquirer(@NonNull String tag) {
-            Objects.requireNonNull(tag);
-            return new SleepTokenAcquirerImpl(tag);
-        }
 
         @Override
         public ComponentName getHomeActivityForUser(int userId) {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index e8a3951..10e0641 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -735,8 +735,6 @@
 
     /** All tokens used to put activities on this root task to sleep (including mOffToken) */
     final ArrayList<RootWindowContainer.SleepToken> mAllSleepTokens = new ArrayList<>();
-    /** The token acquirer to put root tasks on the display to sleep */
-    private final ActivityTaskManagerInternal.SleepTokenAcquirer mOffTokenAcquirer;
 
     private boolean mSleeping;
 
@@ -1131,7 +1129,6 @@
         mDisplay = display;
         mDisplayId = display.getDisplayId();
         mCurrentUniqueDisplayId = display.getUniqueId();
-        mOffTokenAcquirer = mRootWindowContainer.mDisplayOffTokenAcquirer;
         mWallpaperController = new WallpaperController(mWmService, this);
         mWallpaperController.resetLargestDisplay(display);
         display.getDisplayInfo(mDisplayInfo);
@@ -6157,9 +6154,9 @@
         final int displayState = mDisplayInfo.state;
         if (displayId != DEFAULT_DISPLAY) {
             if (displayState == Display.STATE_OFF) {
-                mOffTokenAcquirer.acquire(mDisplayId);
+                mRootWindowContainer.mDisplayOffTokenAcquirer.acquire(mDisplayId);
             } else if (displayState == Display.STATE_ON) {
-                mOffTokenAcquirer.release(mDisplayId);
+                mRootWindowContainer.mDisplayOffTokenAcquirer.release(mDisplayId);
             }
             ProtoLog.v(WM_DEBUG_CONTENT_RECORDING,
                     "Content Recording: Display %d state was (%d), is now (%d), so update "
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index 5c62120..107d31e 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -804,6 +804,14 @@
                     mAwake /* waiting */);
             if (!awake) {
                 onDisplaySwitchFinished();
+                // In case PhoneWindowManager's startedGoingToSleep is called after screenTurnedOff
+                // (the source caller is in order but the methods run on different threads) and
+                // updateScreenOffSleepToken was skipped by mIsGoingToSleepDefaultDisplay. Then
+                // acquire sleep token if screen is off.
+                if (!mScreenOnEarly && !mScreenOnFully && !mDisplayContent.isSleeping()) {
+                    Slog.w(TAG, "Late acquire sleep token for " + mDisplayContent);
+                    mService.mRoot.mDisplayOffTokenAcquirer.acquire(mDisplayContent.mDisplayId);
+                }
             }
         }
     }
@@ -851,6 +859,7 @@
     public void screenTurningOn(ScreenOnListener screenOnListener) {
         WindowProcessController visibleDozeUiProcess = null;
         synchronized (mLock) {
+            mService.mRoot.mDisplayOffTokenAcquirer.release(mDisplayContent.mDisplayId);
             mScreenOnEarly = true;
             mScreenOnFully = false;
             mKeyguardDrawComplete = false;
@@ -875,8 +884,12 @@
         onDisplaySwitchFinished();
     }
 
-    public void screenTurnedOff() {
+    /** It is called after {@link #screenTurningOn}. This runs on PowerManager's thread. */
+    public void screenTurnedOff(boolean acquireSleepToken) {
         synchronized (mLock) {
+            if (acquireSleepToken) {
+                mService.mRoot.mDisplayOffTokenAcquirer.acquire(mDisplayContent.mDisplayId);
+            }
             mScreenOnEarly = false;
             mScreenOnFully = false;
             mKeyguardDrawComplete = false;
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index 5d8a96c..0c489d6 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -87,7 +87,7 @@
     private final SparseArray<KeyguardDisplayState> mDisplayStates = new SparseArray<>();
     private final ActivityTaskManagerService mService;
     private RootWindowContainer mRootWindowContainer;
-    private final ActivityTaskManagerInternal.SleepTokenAcquirer mSleepTokenAcquirer;
+    private final ActivityTaskManagerService.SleepTokenAcquirer mSleepTokenAcquirer;
     private boolean mWaitingForWakeTransition;
     private Transition.ReadyCondition mWaitAodHide = null;
 
@@ -95,7 +95,7 @@
             ActivityTaskSupervisor taskSupervisor) {
         mService = service;
         mTaskSupervisor = taskSupervisor;
-        mSleepTokenAcquirer = mService.new SleepTokenAcquirerImpl(KEYGUARD_SLEEP_TOKEN_TAG);
+        mSleepTokenAcquirer = mService.new SleepTokenAcquirer(KEYGUARD_SLEEP_TOKEN_TAG);
     }
 
     void setWindowManager(WindowManagerService windowManager) {
@@ -658,10 +658,10 @@
 
         private boolean mRequestDismissKeyguard;
         private final ActivityTaskManagerService mService;
-        private final ActivityTaskManagerInternal.SleepTokenAcquirer mSleepTokenAcquirer;
+        private final ActivityTaskManagerService.SleepTokenAcquirer mSleepTokenAcquirer;
 
         KeyguardDisplayState(ActivityTaskManagerService service, int displayId,
-                ActivityTaskManagerInternal.SleepTokenAcquirer acquirer) {
+                ActivityTaskManagerService.SleepTokenAcquirer acquirer) {
             mService = service;
             mDisplayId = displayId;
             mSleepTokenAcquirer = acquirer;
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 866dcd5..8f5612c 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -215,7 +215,7 @@
     private static final String DISPLAY_OFF_SLEEP_TOKEN_TAG = "Display-off";
 
     /** The token acquirer to put root tasks on the displays to sleep */
-    final ActivityTaskManagerInternal.SleepTokenAcquirer mDisplayOffTokenAcquirer;
+    final ActivityTaskManagerService.SleepTokenAcquirer mDisplayOffTokenAcquirer;
 
     /**
      * The modes which affect which tasks are returned when calling
@@ -450,7 +450,7 @@
         mService = service.mAtmService;
         mTaskSupervisor = mService.mTaskSupervisor;
         mTaskSupervisor.mRootWindowContainer = this;
-        mDisplayOffTokenAcquirer = mService.new SleepTokenAcquirerImpl(DISPLAY_OFF_SLEEP_TOKEN_TAG);
+        mDisplayOffTokenAcquirer = mService.new SleepTokenAcquirer(DISPLAY_OFF_SLEEP_TOKEN_TAG);
         mDeviceStateController = new DeviceStateController(service.mContext, service.mGlobalLock);
         mDisplayRotationCoordinator = new DisplayRotationCoordinator();
     }
diff --git a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
index 5aa34d2..92953e5 100644
--- a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
@@ -17,6 +17,8 @@
 package com.android.server.wm;
 
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENTS_INFO;
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO;
 import static android.window.TaskFragmentOrganizer.putErrorInfoInBundle;
 import static android.window.TaskFragmentTransaction.TYPE_ACTIVITY_REPARENTED_TO_TASK;
 import static android.window.TaskFragmentTransaction.TYPE_TASK_FRAGMENT_APPEARED;
@@ -206,7 +208,13 @@
             mOrganizerPid = pid;
             mAppThread = getAppThread(pid, mOrganizerUid);
             for (int i = mOrganizedTaskFragments.size() - 1; i >= 0; i--) {
-                mOrganizedTaskFragments.get(i).onTaskFragmentOrganizerRestarted(organizer);
+                final TaskFragment taskFragment = mOrganizedTaskFragments.get(i);
+                if (taskFragment.isAttached()
+                        && taskFragment.getTopNonFinishingActivity() != null) {
+                    taskFragment.onTaskFragmentOrganizerRestarted(organizer);
+                } else {
+                    mOrganizedTaskFragments.remove(taskFragment);
+                }
             }
             try {
                 mOrganizer.asBinder().linkToDeath(this, 0 /*flags*/);
@@ -575,8 +583,29 @@
         }
 
         mCachedTaskFragmentOrganizerStates.remove(cachedState);
-        outSavedState.putAll(cachedState.mSavedState);
         cachedState.restore(organizer, pid);
+        outSavedState.putAll(cachedState.mSavedState);
+
+        // Collect the organized TfInfo and TfParentInfo in the system.
+        final ArrayList<TaskFragmentInfo> infos = new ArrayList<>();
+        final ArrayMap<Integer, Task> tasks = new ArrayMap<>();
+        final int fragmentCount = cachedState.mOrganizedTaskFragments.size();
+        for (int j = 0; j < fragmentCount; j++) {
+            final TaskFragment tf = cachedState.mOrganizedTaskFragments.get(j);
+            infos.add(tf.getTaskFragmentInfo());
+            if (!tasks.containsKey(tf.getTask().mTaskId)) {
+                tasks.put(tf.getTask().mTaskId, tf.getTask());
+            }
+        }
+        outSavedState.putParcelableArrayList(KEY_RESTORE_TASK_FRAGMENTS_INFO, infos);
+
+        final ArrayList<TaskFragmentParentInfo> parentInfos = new ArrayList<>();
+        for (int j = tasks.size() - 1; j >= 0; j--) {
+            parentInfos.add(tasks.valueAt(j).getTaskFragmentParentInfo());
+        }
+        outSavedState.putParcelableArrayList(KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO,
+                parentInfos);
+
         mTaskFragmentOrganizerState.put(organizer.asBinder(), cachedState);
         mPendingTaskFragmentEvents.put(organizer.asBinder(), new ArrayList<>());
         return true;
diff --git a/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTests.java b/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTests.java
index e694c0b..536dcfb 100644
--- a/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTests.java
+++ b/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTests.java
@@ -42,7 +42,6 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.clearInvocations;
 
 import android.app.ActivityManager;
 import android.app.AppOpsManager;
@@ -135,15 +134,13 @@
         doNothing().when(mPhoneWindowManager).initializeHdmiState();
         final boolean[] isScreenTurnedOff = { false };
         final DisplayPolicy displayPolicy = mock(DisplayPolicy.class);
-        doAnswer(invocation -> isScreenTurnedOff[0] = true).when(displayPolicy).screenTurnedOff();
+        doAnswer(invocation -> isScreenTurnedOff[0] = true).when(displayPolicy).screenTurnedOff(
+                anyBoolean());
         doAnswer(invocation -> !isScreenTurnedOff[0]).when(displayPolicy).isScreenOnEarly();
         doAnswer(invocation -> !isScreenTurnedOff[0]).when(displayPolicy).isScreenOnFully();
 
         mPhoneWindowManager.mDefaultDisplayPolicy = displayPolicy;
         mPhoneWindowManager.mDefaultDisplayRotation = mock(DisplayRotation.class);
-        final ActivityTaskManagerInternal.SleepTokenAcquirer tokenAcquirer =
-                mock(ActivityTaskManagerInternal.SleepTokenAcquirer.class);
-        doReturn(tokenAcquirer).when(mAtmInternal).createSleepTokenAcquirer(anyString());
         final PowerManager pm = mock(PowerManager.class);
         doReturn(true).when(pm).isInteractive();
         doReturn(pm).when(mContext).getSystemService(eq(Context.POWER_SERVICE));
@@ -155,9 +152,8 @@
         assertThat(mPhoneWindowManager.mIsGoingToSleepDefaultDisplay).isFalse();
 
         // Skip sleep-token for non-sleep-screen-off.
-        clearInvocations(tokenAcquirer);
         mPhoneWindowManager.screenTurnedOff(DEFAULT_DISPLAY, true /* isSwappingDisplay */);
-        verify(tokenAcquirer, never()).acquire(anyInt());
+        verify(displayPolicy).screenTurnedOff(false /* acquireSleepToken */);
         assertThat(isScreenTurnedOff[0]).isTrue();
 
         // Apply sleep-token for sleep-screen-off.
@@ -165,21 +161,10 @@
         mPhoneWindowManager.startedGoingToSleep(DEFAULT_DISPLAY, 0 /* reason */);
         assertThat(mPhoneWindowManager.mIsGoingToSleepDefaultDisplay).isTrue();
         mPhoneWindowManager.screenTurnedOff(DEFAULT_DISPLAY, true /* isSwappingDisplay */);
-        verify(tokenAcquirer).acquire(eq(DEFAULT_DISPLAY));
+        verify(displayPolicy).screenTurnedOff(true /* acquireSleepToken */);
 
         mPhoneWindowManager.finishedGoingToSleep(DEFAULT_DISPLAY, 0 /* reason */);
         assertThat(mPhoneWindowManager.mIsGoingToSleepDefaultDisplay).isFalse();
-
-        // Simulate unexpected reversed order: screenTurnedOff -> startedGoingToSleep. The sleep
-        // token can still be acquired.
-        isScreenTurnedOff[0] = false;
-        clearInvocations(tokenAcquirer);
-        mPhoneWindowManager.screenTurnedOff(DEFAULT_DISPLAY, true /* isSwappingDisplay */);
-        verify(tokenAcquirer, never()).acquire(anyInt());
-        assertThat(displayPolicy.isScreenOnEarly()).isFalse();
-        assertThat(displayPolicy.isScreenOnFully()).isFalse();
-        mPhoneWindowManager.startedGoingToSleep(DEFAULT_DISPLAY, 0 /* reason */);
-        verify(tokenAcquirer).acquire(eq(DEFAULT_DISPLAY));
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 1e035da..e2e76d6 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -3231,7 +3231,7 @@
         mDisplayContent.mOpeningApps.remove(activity);
         mDisplayContent.mClosingApps.remove(activity);
         activity.commitVisibility(false /* visible */, false /* performLayout */);
-        mDisplayContent.getDisplayPolicy().screenTurnedOff();
+        mDisplayContent.getDisplayPolicy().screenTurnedOff(false /* acquireSleepToken */);
         final KeyguardController controller = mSupervisor.getKeyguardController();
         doReturn(true).when(controller).isKeyguardGoingAway(anyInt());
         activity.setVisibility(true);
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
index caeb41c..f32a234 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
@@ -284,11 +284,11 @@
         final DisplayPolicy policy = mDisplayContent.getDisplayPolicy();
         policy.addWindowLw(mNotificationShadeWindow, mNotificationShadeWindow.mAttrs);
 
-        policy.screenTurnedOff();
+        policy.screenTurnedOff(false /* acquireSleepToken */);
         policy.setAwake(false);
         policy.screenTurningOn(null /* screenOnListener */);
         assertTrue(wpc.isShowingUiWhileDozing());
-        policy.screenTurnedOff();
+        policy.screenTurnedOff(false /* acquireSleepToken */);
         assertFalse(wpc.isShowingUiWhileDozing());
 
         policy.screenTurningOn(null /* screenOnListener */);
@@ -393,7 +393,7 @@
                 info.logicalWidth, info.logicalHeight).mConfigFrame);
 
         // If screen is not fully turned on, then the cache should be preserved.
-        displayPolicy.screenTurnedOff();
+        displayPolicy.screenTurnedOff(false /* acquireSleepToken */);
         final TransitionController transitionController = mDisplayContent.mTransitionController;
         spyOn(transitionController);
         doReturn(true).when(transitionController).isCollecting();
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
index cc1805a..fd959b9 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
@@ -242,7 +242,7 @@
 
         final Rect relStartBounds = new Rect(mTaskFragment.getRelativeEmbeddedBounds());
         final DisplayPolicy displayPolicy = mDisplayContent.getDisplayPolicy();
-        displayPolicy.screenTurnedOff();
+        displayPolicy.screenTurnedOff(false /* acquireSleepToken */);
 
         assertFalse(mTaskFragment.okToAnimate());