Fix snapshot starting window stuck if the task never gain focus

The tasksnapshot starting window could stuck on task if the task never
gain focus, this could happen when launch multiple tasks to front, e.g.
split screen, or add snapshot starting window to a pip task.

As using task focus to judge the task snapshot removal may error-prone.
(i.e. unexpected window focus in/lost, non-focusable task window..)

To ensure the tasksnapshot removal works stable, extanding
MAX_DELAY_REMOVAL_TIME_IME_VISIBLE timeout from 450ms to 600ms and also
reference mayImeShowOnLaunchingActivity to know whether IME showing on
this activity.

Bug: 199377815
Bug: 201264769
Bug: 200778734
Test: atest ActivityRecordTests StartingSurfaceDrawerTests
SplashscreenTests
Test: manual launch apps with IME from Recents.
Test: enter pip, power on/off, verify starting window is removed.
Test: manual enter split screen, verify starting window is removed.

Change-Id: I81b048a655124923a5f19e7f236511502bbf4c91
(cherry picked from commit a4f760bee7d86d8d68d155dc256f984f1b3d7478)

18 files changed, 205 insertions, 109 deletions

diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 887f7617a17f..ceab0560e0d3 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -3191,13 +3191,6 @@ package android.window {
     method @Nullable public android.view.View getBrandingView();
   }
 
-  public final class StartingWindowInfo implements android.os.Parcelable {
-    ctor public StartingWindowInfo();
-    method public int describeContents();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.window.StartingWindowInfo> CREATOR;
-  }
-
   public final class TaskAppearedInfo implements android.os.Parcelable {
     ctor public TaskAppearedInfo(@NonNull android.app.ActivityManager.RunningTaskInfo, @NonNull android.view.SurfaceControl);
     method public int describeContents();
@@ -3264,7 +3257,6 @@ package android.window {
 
   public class TaskOrganizer extends android.window.WindowOrganizer {
     ctor public TaskOrganizer();
-    method @BinderThread public void addStartingWindow(@NonNull android.window.StartingWindowInfo, @NonNull android.os.IBinder);
     method @BinderThread public void copySplashScreenView(int);
     method @Nullable @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS) public void createRootTask(int, int, @Nullable android.os.IBinder);
     method @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS) public boolean deleteRootTask(@NonNull android.window.WindowContainerToken);
@@ -3277,7 +3269,6 @@ package android.window {
     method @BinderThread public void onTaskInfoChanged(@NonNull android.app.ActivityManager.RunningTaskInfo);
     method @BinderThread public void onTaskVanished(@NonNull android.app.ActivityManager.RunningTaskInfo);
     method @CallSuper @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS) public java.util.List<android.window.TaskAppearedInfo> registerOrganizer();
-    method @BinderThread public void removeStartingWindow(int, @Nullable android.view.SurfaceControl, @Nullable android.graphics.Rect, boolean);
     method @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS) public void setInterceptBackPressedOnTaskRoot(@NonNull android.window.WindowContainerToken, boolean);
     method @CallSuper @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS) public void unregisterOrganizer();
   }
diff --git a/core/java/android/window/ITaskOrganizer.aidl b/core/java/android/window/ITaskOrganizer.aidl
index 69bc1b5f7763..fd86769293a6 100644
--- a/core/java/android/window/ITaskOrganizer.aidl
+++ b/core/java/android/window/ITaskOrganizer.aidl
@@ -20,6 +20,7 @@ import android.view.SurfaceControl;
 import android.app.ActivityManager;
 import android.graphics.Rect;
 import android.window.StartingWindowInfo;
+import android.window.StartingWindowRemovalInfo;
 import android.window.WindowContainerToken;
 
 /**
@@ -39,12 +40,9 @@ oneway interface ITaskOrganizer {
 
     /**
      * Called when the Task want to remove the starting window.
-     * @param leash A persistent leash for the top window in this task.
-     * @param frame Window frame of the top window.
-     * @param playRevealAnimation Play vanish animation.
+     * @param removalInfo The information used to remove the starting window.
      */
-    void removeStartingWindow(int taskId, in SurfaceControl leash, in Rect frame,
-            in boolean playRevealAnimation);
+    void removeStartingWindow(in StartingWindowRemovalInfo removalInfo);
 
     /**
      * Called when the Task want to copy the splash screen.
diff --git a/core/java/android/window/StartingWindowInfo.java b/core/java/android/window/StartingWindowInfo.java
index 10d21a0ff003..5950e9fc6456 100644
--- a/core/java/android/window/StartingWindowInfo.java
+++ b/core/java/android/window/StartingWindowInfo.java
@@ -19,7 +19,6 @@ package android.window;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.annotation.TestApi;
 import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.content.pm.ActivityInfo;
@@ -34,7 +33,6 @@ import android.view.WindowManager;
  * start in the system.
  * @hide
  */
-@TestApi
 public final class StartingWindowInfo implements Parcelable {
     /**
      * Prefer nothing or not care the type of starting window.
diff --git a/core/java/android/window/StartingWindowRemovalInfo.aidl b/core/java/android/window/StartingWindowRemovalInfo.aidl
new file mode 100644
index 000000000000..8e4ac0453f83
--- /dev/null
+++ b/core/java/android/window/StartingWindowRemovalInfo.aidl
@@ -0,0 +1,20 @@
+/**
+ * Copyright (c) 2021, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.window;
+
+/** @hide */
+parcelable StartingWindowRemovalInfo;
\ No newline at end of file
diff --git a/core/java/android/window/StartingWindowRemovalInfo.java b/core/java/android/window/StartingWindowRemovalInfo.java
new file mode 100644
index 000000000000..573db0d58625
--- /dev/null
+++ b/core/java/android/window/StartingWindowRemovalInfo.java
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.window;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.graphics.Rect;
+import android.os.Parcel;
+import android.os.Parcelable;
+import android.view.SurfaceControl;
+
+/**
+ * Information when removing a starting window of a particular task.
+ * @hide
+ */
+public final class StartingWindowRemovalInfo implements Parcelable {
+
+    /**
+     * The identifier of a task.
+     * @hide
+     */
+    public int taskId;
+
+    /**
+     * The animation container layer of the top activity.
+     * @hide
+     */
+    @Nullable
+    public SurfaceControl windowAnimationLeash;
+
+    /**
+     * The main window frame for the window of the top activity.
+     * @hide
+     */
+    @Nullable
+    public Rect mainFrame;
+
+    /**
+     * Whether need to play reveal animation.
+     * @hide
+     */
+    public boolean playRevealAnimation;
+
+    /**
+     * Whether need to defer removing the starting window for IME.
+     * @hide
+     */
+    public boolean deferRemoveForIme;
+
+    public StartingWindowRemovalInfo() {
+
+    }
+
+    private StartingWindowRemovalInfo(@NonNull Parcel source) {
+        readFromParcel(source);
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    void readFromParcel(@NonNull Parcel source) {
+        taskId = source.readInt();
+        windowAnimationLeash = source.readTypedObject(SurfaceControl.CREATOR);
+        mainFrame = source.readTypedObject(Rect.CREATOR);
+        playRevealAnimation = source.readBoolean();
+        deferRemoveForIme = source.readBoolean();
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeInt(taskId);
+        dest.writeTypedObject(windowAnimationLeash, flags);
+        dest.writeTypedObject(mainFrame, flags);
+        dest.writeBoolean(playRevealAnimation);
+        dest.writeBoolean(deferRemoveForIme);
+    }
+
+    @Override
+    public String toString() {
+        return "StartingWindowRemovalInfo{taskId=" + taskId
+                + " frame=" + mainFrame
+                + " playRevealAnimation=" + playRevealAnimation
+                + " deferRemoveForIme=" + deferRemoveForIme + "}";
+    }
+
+    public static final @android.annotation.NonNull Creator<StartingWindowRemovalInfo> CREATOR =
+            new Creator<StartingWindowRemovalInfo>() {
+                public StartingWindowRemovalInfo createFromParcel(@NonNull Parcel source) {
+                    return new StartingWindowRemovalInfo(source);
+                }
+                public StartingWindowRemovalInfo[] newArray(int size) {
+                    return new StartingWindowRemovalInfo[size];
+                }
+            };
+}
diff --git a/core/java/android/window/TaskOrganizer.java b/core/java/android/window/TaskOrganizer.java
index 845c13d9a01b..27c7d3158f95 100644
--- a/core/java/android/window/TaskOrganizer.java
+++ b/core/java/android/window/TaskOrganizer.java
@@ -24,7 +24,6 @@ import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.annotation.TestApi;
 import android.app.ActivityManager;
-import android.graphics.Rect;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.view.SurfaceControl;
@@ -94,6 +93,7 @@ public class TaskOrganizer extends WindowOrganizer {
      * @param info The information about the Task that's available
      * @param appToken Token of the application being started.
      *        context to for resources
+     * @hide
      */
     @BinderThread
     public void addStartingWindow(@NonNull StartingWindowInfo info,
@@ -101,14 +101,11 @@ public class TaskOrganizer extends WindowOrganizer {
 
     /**
      * Called when the Task want to remove the starting window.
-     * @param leash A persistent leash for the top window in this task. Release it once exit
-     *              animation has finished.
-     * @param frame Window frame of the top window.
-     * @param playRevealAnimation Play vanish animation.
+     * @param removalInfo The information used to remove the starting window.
+     * @hide
      */
     @BinderThread
-    public void removeStartingWindow(int taskId, @Nullable SurfaceControl leash,
-            @Nullable Rect frame, boolean playRevealAnimation) {}
+    public void removeStartingWindow(@NonNull StartingWindowRemovalInfo removalInfo) {}
 
     /**
      * Called when the Task want to copy the splash screen.
@@ -257,10 +254,8 @@ public class TaskOrganizer extends WindowOrganizer {
         }
 
         @Override
-        public void removeStartingWindow(int taskId, SurfaceControl leash, Rect frame,
-                boolean playRevealAnimation) {
-            mExecutor.execute(() -> TaskOrganizer.this.removeStartingWindow(taskId, leash, frame,
-                    playRevealAnimation));
+        public void removeStartingWindow(StartingWindowRemovalInfo removalInfo) {
+            mExecutor.execute(() -> TaskOrganizer.this.removeStartingWindow(removalInfo));
         }
 
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index d925a9218a77..020ecb7186ed 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -42,6 +42,7 @@ import android.util.SparseArray;
 import android.view.SurfaceControl;
 import android.window.ITaskOrganizerController;
 import android.window.StartingWindowInfo;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskAppearedInfo;
 import android.window.TaskOrganizer;
 
@@ -322,10 +323,9 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
     }
 
     @Override
-    public void removeStartingWindow(int taskId, SurfaceControl leash, Rect frame,
-            boolean playRevealAnimation) {
+    public void removeStartingWindow(StartingWindowRemovalInfo removalInfo) {
         if (mStartingWindow != null) {
-            mStartingWindow.removeStartingWindow(taskId, leash, frame, playRevealAnimation);
+            mStartingWindow.removeStartingWindow(removalInfo);
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
index ff3428cd9a6e..979bf0056b72 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
@@ -37,7 +37,6 @@ import android.content.res.Resources;
 import android.content.res.TypedArray;
 import android.graphics.Color;
 import android.graphics.PixelFormat;
-import android.graphics.Rect;
 import android.hardware.display.DisplayManager;
 import android.os.IBinder;
 import android.os.RemoteCallback;
@@ -48,7 +47,6 @@ import android.util.Slog;
 import android.util.SparseArray;
 import android.view.Choreographer;
 import android.view.Display;
-import android.view.SurfaceControl;
 import android.view.SurfaceControlViewHost;
 import android.view.View;
 import android.view.WindowManager;
@@ -58,6 +56,7 @@ import android.window.SplashScreenView;
 import android.window.SplashScreenView.SplashScreenViewParcelable;
 import android.window.StartingWindowInfo;
 import android.window.StartingWindowInfo.StartingWindowType;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskSnapshot;
 
 import com.android.internal.R;
@@ -118,6 +117,7 @@ public class StartingSurfaceDrawer {
     private Choreographer mChoreographer;
     private final WindowManagerGlobal mWindowManagerGlobal;
     private StartingSurface.SysuiProxy mSysuiProxy;
+    private final StartingWindowRemovalInfo mTmpRemovalInfo = new StartingWindowRemovalInfo();
 
     /**
      * @param splashScreenExecutor The thread used to control add and remove starting window.
@@ -458,12 +458,13 @@ public class StartingSurfaceDrawer {
     /**
      * Called when the content of a task is ready to show, starting window can be removed.
      */
-    public void removeStartingWindow(int taskId, SurfaceControl leash, Rect frame,
-            boolean playRevealAnimation) {
+    public void removeStartingWindow(StartingWindowRemovalInfo removalInfo) {
         if (DEBUG_SPLASH_SCREEN || DEBUG_TASK_SNAPSHOT) {
-            Slog.d(TAG, "Task start finish, remove starting surface for task " + taskId);
+            Slog.d(TAG, "Task start finish, remove starting surface for task "
+                    + removalInfo.taskId);
         }
-        removeWindowSynced(taskId, leash, frame, playRevealAnimation);
+        removeWindowSynced(removalInfo);
+
     }
 
     /**
@@ -556,7 +557,8 @@ public class StartingSurfaceDrawer {
     }
 
     private void removeWindowNoAnimate(int taskId) {
-        removeWindowSynced(taskId, null, null, false);
+        mTmpRemovalInfo.taskId = taskId;
+        removeWindowSynced(mTmpRemovalInfo);
     }
 
     void onImeDrawnOnTask(int taskId) {
@@ -568,8 +570,8 @@ public class StartingSurfaceDrawer {
         mStartingWindowRecords.remove(taskId);
     }
 
-    protected void removeWindowSynced(int taskId, SurfaceControl leash, Rect frame,
-            boolean playRevealAnimation) {
+    protected void removeWindowSynced(StartingWindowRemovalInfo removalInfo) {
+        final int taskId = removalInfo.taskId;
         final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
         if (record != null) {
             if (record.mDecorView != null) {
@@ -580,9 +582,9 @@ public class StartingSurfaceDrawer {
                     if (record.mSuggestType == STARTING_WINDOW_TYPE_LEGACY_SPLASH_SCREEN) {
                         removeWindowInner(record.mDecorView, false);
                     } else {
-                        if (playRevealAnimation) {
+                        if (removalInfo.playRevealAnimation) {
                             mSplashscreenContentDrawer.applyExitAnimation(record.mContentView,
-                                    leash, frame,
+                                    removalInfo.windowAnimationLeash, removalInfo.mainFrame,
                                     () -> removeWindowInner(record.mDecorView, true));
                         } else {
                             // the SplashScreenView has been copied to client, hide the view to skip
@@ -602,7 +604,7 @@ public class StartingSurfaceDrawer {
                     Slog.v(TAG, "Removing task snapshot window for " + taskId);
                 }
                 record.mTaskSnapshotWindow.scheduleRemove(
-                        () -> mStartingWindowRecords.remove(taskId));
+                        () -> mStartingWindowRecords.remove(taskId), removalInfo.deferRemoveForIme);
             }
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
index 4433e275cd1b..99644f9493d2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
@@ -28,16 +28,14 @@ import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
 import android.content.Context;
 import android.graphics.Color;
-import android.graphics.Rect;
-import android.os.Build;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.Trace;
 import android.util.Slog;
 import android.util.SparseIntArray;
-import android.view.SurfaceControl;
 import android.window.StartingWindowInfo;
 import android.window.StartingWindowInfo.StartingWindowType;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskOrganizer;
 import android.window.TaskSnapshot;
 
@@ -68,7 +66,7 @@ import com.android.wm.shell.common.TransactionPool;
 public class StartingWindowController implements RemoteCallable<StartingWindowController> {
     private static final String TAG = StartingWindowController.class.getSimpleName();
 
-    public static final boolean DEBUG_SPLASH_SCREEN = Build.isDebuggable();
+    public static final boolean DEBUG_SPLASH_SCREEN = false;
     public static final boolean DEBUG_TASK_SNAPSHOT = false;
 
     private static final long TASK_BG_COLOR_RETAIN_TIME_MS = 5000;
@@ -186,13 +184,12 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
     /**
      * Called when the content of a task is ready to show, starting window can be removed.
      */
-    public void removeStartingWindow(int taskId, SurfaceControl leash, Rect frame,
-            boolean playRevealAnimation) {
+    public void removeStartingWindow(StartingWindowRemovalInfo removalInfo) {
         mSplashScreenExecutor.execute(() -> mStartingSurfaceDrawer.removeStartingWindow(
-                taskId, leash, frame, playRevealAnimation));
+                removalInfo));
         mSplashScreenExecutor.executeDelayed(() -> {
             synchronized (mTaskBackgroundColors) {
-                mTaskBackgroundColors.delete(taskId);
+                mTaskBackgroundColors.delete(removalInfo.taskId);
             }
         }, TASK_BG_COLOR_RETAIN_TIME_MS);
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
index 979aa1f4ac46..3e88c464d359 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
@@ -123,14 +123,13 @@ public class TaskSnapshotWindow {
      * Ideally the delay time will be shorter when receiving
      * {@link StartingSurfaceDrawer#onImeDrawnOnTask(int)}.
      */
-    private static final long MAX_DELAY_REMOVAL_TIME_IME_VISIBLE = 450;
+    private static final long MAX_DELAY_REMOVAL_TIME_IME_VISIBLE = 600;
 
     //tmp vars for unused relayout params
     private static final Point TMP_SURFACE_SIZE = new Point();
 
     private final Window mWindow;
     private final Runnable mClearWindowHandler;
-    private final long mDelayRemovalTime;
     private final ShellExecutor mSplashScreenExecutor;
     private final SurfaceControl mSurfaceControl;
     private final IWindowSession mSession;
@@ -221,13 +220,10 @@ public class TaskSnapshotWindow {
             taskDescription.setBackgroundColor(WHITE);
         }
 
-        final long delayRemovalTime = snapshot.hasImeSurface() ? MAX_DELAY_REMOVAL_TIME_IME_VISIBLE
-                : DELAY_REMOVAL_TIME_GENERAL;
-
         final TaskSnapshotWindow snapshotSurface = new TaskSnapshotWindow(
                 surfaceControl, snapshot, layoutParams.getTitle(), taskDescription, appearance,
                 windowFlags, windowPrivateFlags, taskBounds, orientation, activityType,
-                delayRemovalTime, topWindowInsetsState, clearWindowHandler, splashScreenExecutor);
+                topWindowInsetsState, clearWindowHandler, splashScreenExecutor);
         final Window window = snapshotSurface.mWindow;
 
         final InsetsState tmpInsetsState = new InsetsState();
@@ -265,9 +261,8 @@ public class TaskSnapshotWindow {
     public TaskSnapshotWindow(SurfaceControl surfaceControl,
             TaskSnapshot snapshot, CharSequence title, TaskDescription taskDescription,
             int appearance, int windowFlags, int windowPrivateFlags, Rect taskBounds,
-            int currentOrientation, int activityType, long delayRemovalTime,
-            InsetsState topWindowInsetsState, Runnable clearWindowHandler,
-            ShellExecutor splashScreenExecutor) {
+            int currentOrientation, int activityType, InsetsState topWindowInsetsState,
+            Runnable clearWindowHandler, ShellExecutor splashScreenExecutor) {
         mSplashScreenExecutor = splashScreenExecutor;
         mSession = WindowManagerGlobal.getWindowSession();
         mWindow = new Window();
@@ -283,7 +278,6 @@ public class TaskSnapshotWindow {
         mStatusBarColor = taskDescription.getStatusBarColor();
         mOrientationOnCreation = currentOrientation;
         mActivityType = activityType;
-        mDelayRemovalTime = delayRemovalTime;
         mTransaction = new SurfaceControl.Transaction();
         mClearWindowHandler = clearWindowHandler;
         mHasImeSurface = snapshot.hasImeSurface();
@@ -314,7 +308,7 @@ public class TaskSnapshotWindow {
         mSystemBarBackgroundPainter.drawNavigationBarBackground(c);
     }
 
-    void scheduleRemove(Runnable onRemove) {
+    void scheduleRemove(Runnable onRemove, boolean deferRemoveForIme) {
         // Show the latest content as soon as possible for unlocking to home.
         if (mActivityType == ACTIVITY_TYPE_HOME) {
             removeImmediately();
@@ -329,9 +323,12 @@ public class TaskSnapshotWindow {
             TaskSnapshotWindow.this.removeImmediately();
             onRemove.run();
         };
-        mSplashScreenExecutor.executeDelayed(mScheduledRunnable, mDelayRemovalTime);
+        final long delayRemovalTime = mHasImeSurface && deferRemoveForIme
+                ? MAX_DELAY_REMOVAL_TIME_IME_VISIBLE
+                : DELAY_REMOVAL_TIME_GENERAL;
+        mSplashScreenExecutor.executeDelayed(mScheduledRunnable, delayRemovalTime);
         if (DEBUG) {
-            Slog.d(TAG, "Defer removing snapshot surface in " + mDelayRemovalTime);
+            Slog.d(TAG, "Defer removing snapshot surface in " + delayRemovalTime);
         }
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
index 36722d9147ab..e5a8aa043d1a 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
@@ -57,12 +57,12 @@ import android.view.Display;
 import android.view.IWindowSession;
 import android.view.InsetsState;
 import android.view.Surface;
-import android.view.SurfaceControl;
 import android.view.View;
 import android.view.WindowManager;
 import android.view.WindowManagerGlobal;
 import android.view.WindowMetrics;
 import android.window.StartingWindowInfo;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskSnapshot;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -119,10 +119,9 @@ public class StartingSurfaceDrawerTests {
         }
 
         @Override
-        protected void removeWindowSynced(int taskId, SurfaceControl leash, Rect frame,
-                boolean playRevealAnimation) {
+        protected void removeWindowSynced(StartingWindowRemovalInfo removalInfo) {
             // listen for removeView
-            if (mAddWindowForTask == taskId) {
+            if (mAddWindowForTask == removalInfo.taskId) {
                 mAddWindowForTask = 0;
             }
         }
@@ -172,9 +171,11 @@ public class StartingSurfaceDrawerTests {
                 eq(STARTING_WINDOW_TYPE_SPLASH_SCREEN));
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, taskId);
 
-        mStartingSurfaceDrawer.removeStartingWindow(windowInfo.taskInfo.taskId, null, null, false);
+        StartingWindowRemovalInfo removalInfo = new StartingWindowRemovalInfo();
+        removalInfo.taskId = windowInfo.taskInfo.taskId;
+        mStartingSurfaceDrawer.removeStartingWindow(removalInfo);
         waitHandlerIdle(mTestHandler);
-        verify(mStartingSurfaceDrawer).removeWindowSynced(eq(taskId), any(), any(), eq(false));
+        verify(mStartingSurfaceDrawer).removeWindowSynced(any());
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, 0);
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/TaskSnapshotWindowTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/TaskSnapshotWindowTest.java
index a098a6863493..aad9528bd527 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/TaskSnapshotWindowTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/TaskSnapshotWindowTest.java
@@ -83,8 +83,7 @@ public class TaskSnapshotWindowTest {
                 createTaskDescription(Color.WHITE, Color.RED, Color.BLUE),
                 0 /* appearance */, windowFlags /* windowFlags */, 0 /* privateWindowFlags */,
                 taskBounds, ORIENTATION_PORTRAIT, ACTIVITY_TYPE_STANDARD,
-                100 /* delayRemovalTime */, new InsetsState(),
-                null /* clearWindow */, new TestShellExecutor());
+                new InsetsState(), null /* clearWindow */, new TestShellExecutor());
     }
 
     private TaskSnapshot createTaskSnapshot(int width, int height, Point taskSize,
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 3b1da1d0053d..d1a2e5dd86bf 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -2405,17 +2405,6 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
         });
     }
 
-    void removeStartingWindowIfNeeded() {
-        // Removing the task snapshot after the task is actually focused (see
-        // Task#onWindowFocusChanged). Since some of the app contents may draw in this time and
-        // requires more times to draw finish, in case flicking may happen when removing the task
-        // snapshot too early. (i.e. Showing IME.)
-        if ((mStartingData instanceof SnapshotStartingData) && !getTask().isFocused()) {
-            return;
-        }
-        removeStartingWindow();
-    }
-
     void removeStartingWindow() {
         if (transferSplashScreenIfNeeded()) {
             return;
@@ -6080,13 +6069,13 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
         final Task associatedTask =
                 mSharedStartingData != null ? mSharedStartingData.mAssociatedTask : null;
         if (associatedTask == null) {
-            removeStartingWindowIfNeeded();
+            removeStartingWindow();
         } else if (associatedTask.getActivity(
                 r -> r.mVisibleRequested && !r.firstWindowDrawn) == null) {
             // The last drawn activity may not be the one that owns the starting window.
             final ActivityRecord r = associatedTask.topActivityContainsStartingWindow();
             if (r != null) {
-                r.removeStartingWindowIfNeeded();
+                r.removeStartingWindow();
             }
         }
         updateReportedVisibilityLocked();
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 1db4c1da2023..87e5c94784ca 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -1666,7 +1666,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     }
 
     /** Returns {@code true} if the IME is possible to show on the launching activity. */
-    private boolean mayImeShowOnLaunchingActivity(@NonNull ActivityRecord r) {
+    boolean mayImeShowOnLaunchingActivity(@NonNull ActivityRecord r) {
         final WindowState win = r.findMainWindow();
         if (win == null) {
             return false;
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 0819549ad8df..594b6be489d6 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -4270,7 +4270,7 @@ class Task extends TaskFragment {
     /**
      * @return true if the task is currently focused.
      */
-    boolean isFocused() {
+    private boolean isFocused() {
         if (mDisplayContent == null || mDisplayContent.mFocusedApp == null) {
             return false;
         }
@@ -4328,14 +4328,10 @@ class Task extends TaskFragment {
     }
 
     /**
-     * Called on the task of a window which gained or lost focus.
+     * Called on the task when it gained or lost focus.
      * @param hasFocus
      */
     void onAppFocusChanged(boolean hasFocus) {
-        final ActivityRecord topAct = getTopVisibleActivity();
-        if (topAct != null && (topAct.mStartingData instanceof SnapshotStartingData)) {
-            topAct.removeStartingWindowIfNeeded();
-        }
         updateShadowsRadius(hasFocus, getSyncTransaction());
         dispatchTaskInfoChangedIfNeeded(false /* force */);
     }
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index ccda126b3e95..f1345e8ee61e 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -28,7 +28,6 @@ import android.app.ActivityManager.RunningTaskInfo;
 import android.app.WindowConfiguration;
 import android.content.Intent;
 import android.content.pm.ParceledListSlice;
-import android.graphics.Rect;
 import android.os.Binder;
 import android.os.IBinder;
 import android.os.Parcel;
@@ -40,6 +39,7 @@ import android.window.ITaskOrganizer;
 import android.window.ITaskOrganizerController;
 import android.window.SplashScreenView;
 import android.window.StartingWindowInfo;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskAppearedInfo;
 import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
@@ -498,12 +498,15 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
         if (lastOrganizer == null) {
             return;
         }
-        SurfaceControl windowAnimationLeash = null;
-        Rect mainFrame = null;
+        final StartingWindowRemovalInfo removalInfo = new StartingWindowRemovalInfo();
+        removalInfo.taskId = task.mTaskId;
+        removalInfo.playRevealAnimation = prepareAnimation;
         final boolean playShiftUpAnimation = !task.inMultiWindowMode();
-        if (prepareAnimation && playShiftUpAnimation) {
-            final ActivityRecord topActivity = task.topActivityContainsStartingWindow();
-            if (topActivity != null) {
+        final ActivityRecord topActivity = task.topActivityContainsStartingWindow();
+        if (topActivity != null) {
+            removalInfo.deferRemoveForIme = topActivity.mDisplayContent
+                    .mayImeShowOnLaunchingActivity(topActivity);
+            if (prepareAnimation && playShiftUpAnimation) {
                 final WindowState mainWindow =
                         topActivity.findMainWindow(false/* includeStartingApp */);
                 if (mainWindow != null) {
@@ -512,15 +515,15 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
                     final SurfaceControl.Transaction t = mainWindow.getPendingTransaction();
                     mainWindow.startAnimation(t, adaptor, false,
                             ANIMATION_TYPE_STARTING_REVEAL);
-                    windowAnimationLeash = adaptor.mAnimationLeash;
-                    mainFrame = mainWindow.getRelativeFrame();
-                    t.setPosition(windowAnimationLeash, mainFrame.left, mainFrame.top);
+                    removalInfo.windowAnimationLeash = adaptor.mAnimationLeash;
+                    removalInfo.mainFrame = mainWindow.getRelativeFrame();
+                    t.setPosition(removalInfo.windowAnimationLeash,
+                            removalInfo.mainFrame.left, removalInfo.mainFrame.top);
                 }
             }
         }
         try {
-            lastOrganizer.removeStartingWindow(task.mTaskId, windowAnimationLeash,
-                    mainFrame, prepareAnimation);
+            lastOrganizer.removeStartingWindow(removalInfo);
         } catch (RemoteException e) {
             Slog.e(TAG, "Exception sending onStartTaskFinished callback", e);
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index 17ae2e8f4809..a482bdacfc82 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -80,6 +80,7 @@ import android.view.SurfaceControl;
 import android.window.ITaskOrganizer;
 import android.window.IWindowContainerTransactionCallback;
 import android.window.StartingWindowInfo;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskAppearedInfo;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
@@ -779,8 +780,7 @@ public class WindowOrganizerTests extends WindowTestsBase {
         @Override
         public void addStartingWindow(StartingWindowInfo info, IBinder appToken) { }
         @Override
-        public void removeStartingWindow(int taskId, SurfaceControl leash, Rect frame,
-                boolean playRevealAnimation) { }
+        public void removeStartingWindow(StartingWindowRemovalInfo removalInfo) { }
         @Override
         public void copySplashScreenView(int taskId) { }
         @Override
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index 8ec1bd6c5787..b2d4eea4feae 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -77,7 +77,6 @@ import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
-import android.graphics.Rect;
 import android.hardware.display.DisplayManager;
 import android.os.Build;
 import android.os.Bundle;
@@ -101,6 +100,7 @@ import android.view.WindowManager;
 import android.view.WindowManager.DisplayImePolicy;
 import android.window.ITransitionPlayer;
 import android.window.StartingWindowInfo;
+import android.window.StartingWindowRemovalInfo;
 import android.window.TaskFragmentOrganizer;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
@@ -1457,12 +1457,11 @@ class WindowTestsBase extends SystemServiceTestsBase {
             }
         }
         @Override
-        public void removeStartingWindow(int taskId, SurfaceControl leash, Rect frame,
-                boolean playRevealAnimation) {
+        public void removeStartingWindow(StartingWindowRemovalInfo removalInfo) {
             synchronized (mWMService.mGlobalLock) {
-                final IBinder appToken = mTaskAppMap.get(taskId);
+                final IBinder appToken = mTaskAppMap.get(removalInfo.taskId);
                 if (appToken != null) {
-                    mTaskAppMap.remove(taskId);
+                    mTaskAppMap.remove(removalInfo.taskId);
                     final ActivityRecord activity = mWMService.mRoot.getActivityRecord(
                             appToken);
                     WindowState win = mAppWindowMap.remove(appToken);