RESTRICT AUTOMERGE: More improve IME transition during task switch

This CL aims to optimize the previous CL[1] to schedule removing
tasksnapshot after a fixed timeout according the tasksnapshot:
    - With IME snapshot: 350ms
    - Without IME snapshot: 100ms

As the previous approach has some cons espically when the tasksnapshot
has IME shown:
1) It lacks a signal or callback to notify WmShell to dismiss
tasksnapshot when IME is actually drawn on the task and always
dismissed after the timeout.

2) The timing to schedule tasksnapsit removal is when
ActivityRecord#onWindowFirstDrawn, which is much eariler than the
window focused (about 100-150ms), and it may easier to see flickering
when the task is showing IME.

The reason is that IME is drawn after window focused and started input
connection. Also, starts from R, IME insets visiblity
is handled by the app's UI thread, so if the schedule removal timing
been triggered too early and if IME / App takes more time to handle IME
surface layout, then user might aware the app task flickering when
tasksnapshot dismissed, since IME is not yet be drawn and then it
show up again when the next layout finished.

In this CL, we made the following changes to improve the above cons
- Postpone the schedule removing tasksnapshot (with IME) timing to
  after the app task has focused.
- Modify the tasksnapshot removal timeout (with IME)  from 350ms to
  450ms (with renaming to MAX_DELAY_REMOVAL_TIME_IME_VISIBLE),
  in case some edge cases may take longer time to process IME layout.
- add ITaskOrganizer#onImeDrawnOnTask(taskId) to notify the shell
  task organizer to properly remove the tasksnapshot without waiting
  until the max timeout.

[1]: I7865e17b57961e12a0cdcf068e412195123a6ec7

Fix: 192065018
Test: ateset StartingSurfaceDrawerTests#\
        testRemoveTaskSnapshotWithImeSurfaceWhenOnImeDrawn
Test: manual tests by
   1) launching Android Message with focusing an editor
   2) swiping out to home and launch another apps (e.g. chrome)
   3) swiping up to overview, tapping Android Message task
   4) verify if IME is flickering after switched back.
Change-Id: I81031f64966b1aeb55cc09f381d4d83ec3460dc9
(cherry picked from commit 278e944fd5192b578bc9769323bd46105cddd822)
(cherry picked from commit 5c23301322133f0514cf073cd8c578297b0e1acc)

12 files changed, 204 insertions, 23 deletions

diff --git a/core/java/android/window/ITaskOrganizer.aidl b/core/java/android/window/ITaskOrganizer.aidl
index 8b8dba89ea67..69bc1b5f7763 100644
--- a/core/java/android/window/ITaskOrganizer.aidl
+++ b/core/java/android/window/ITaskOrganizer.aidl
@@ -88,4 +88,9 @@ oneway interface ITaskOrganizer {
      * user has pressed back on the root activity of a task controlled by the task organizer.
      */
     void onBackPressedOnTaskRoot(in ActivityManager.RunningTaskInfo taskInfo);
+
+    /**
+     * Called when the IME has drawn on the organized task.
+     */
+    void onImeDrawnOnTask(int taskId);
 }
diff --git a/core/java/android/window/TaskOrganizer.java b/core/java/android/window/TaskOrganizer.java
index c7c91cdd0941..d8723a821a22 100644
--- a/core/java/android/window/TaskOrganizer.java
+++ b/core/java/android/window/TaskOrganizer.java
@@ -144,6 +144,10 @@ public class TaskOrganizer extends WindowOrganizer {
     @BinderThread
     public void onBackPressedOnTaskRoot(@NonNull ActivityManager.RunningTaskInfo taskInfo) {}
 
+    /** @hide */
+    @BinderThread
+    public void onImeDrawnOnTask(int taskId) {}
+
     /**
      * Creates a persistent root task in WM for a particular windowing-mode.
      * @param displayId The display to create the root task on.
@@ -287,6 +291,11 @@ public class TaskOrganizer extends WindowOrganizer {
         public void onBackPressedOnTaskRoot(ActivityManager.RunningTaskInfo info) {
             mExecutor.execute(() -> TaskOrganizer.this.onBackPressedOnTaskRoot(info));
         }
+
+        @Override
+        public void onImeDrawnOnTask(int taskId) {
+            mExecutor.execute(() -> TaskOrganizer.this.onImeDrawnOnTask(taskId));
+        }
     };
 
     private ITaskOrganizerController getController() {
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 ba0ab6db1003..656bdff0c782 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -336,6 +336,13 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
     }
 
     @Override
+    public void onImeDrawnOnTask(int taskId) {
+        if (mStartingWindow != null) {
+            mStartingWindow.onImeDrawnOnTask(taskId);
+        }
+    }
+
+    @Override
     public void onTaskAppeared(RunningTaskInfo taskInfo, SurfaceControl leash) {
         synchronized (mLock) {
             onTaskAppeared(new TaskAppearedInfo(taskInfo, leash));
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 fc7c86d669cb..147f5e30f9d6 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
@@ -545,6 +545,15 @@ public class StartingSurfaceDrawer {
         removeWindowSynced(taskId, null, null, false);
     }
 
+    void onImeDrawnOnTask(int taskId) {
+        final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
+        if (record != null && record.mTaskSnapshotWindow != null
+                && record.mTaskSnapshotWindow.hasImeSurface()) {
+            record.mTaskSnapshotWindow.removeImmediately();
+        }
+        mStartingWindowRecords.remove(taskId);
+    }
+
     protected void removeWindowSynced(int taskId, SurfaceControl leash, Rect frame,
             boolean playRevealAnimation) {
         final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
@@ -572,14 +581,15 @@ public class StartingSurfaceDrawer {
                     Slog.e(TAG, "Found empty splash screen, remove!");
                     removeWindowInner(record.mDecorView, false);
                 }
+                mStartingWindowRecords.remove(taskId);
             }
             if (record.mTaskSnapshotWindow != null) {
                 if (DEBUG_TASK_SNAPSHOT) {
                     Slog.v(TAG, "Removing task snapshot window for " + taskId);
                 }
-                record.mTaskSnapshotWindow.remove();
+                record.mTaskSnapshotWindow.scheduleRemove(
+                        () -> mStartingWindowRecords.remove(taskId));
             }
-            mStartingWindowRecords.remove(taskId);
         }
     }
 
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 e84d498a9258..dee21b093dce 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
@@ -177,6 +177,13 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
     }
 
     /**
+     * Called when the IME has drawn on the organized task.
+     */
+    public void onImeDrawnOnTask(int taskId) {
+        mSplashScreenExecutor.execute(() -> mStartingSurfaceDrawer.onImeDrawnOnTask(taskId));
+    }
+
+    /**
      * 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,
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 6052d3dee891..2dce20371993 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
@@ -64,7 +64,6 @@ import android.graphics.RectF;
 import android.hardware.HardwareBuffer;
 import android.os.IBinder;
 import android.os.RemoteException;
-import android.os.SystemClock;
 import android.os.Trace;
 import android.util.MergedConfiguration;
 import android.util.Slog;
@@ -119,7 +118,12 @@ public class TaskSnapshotWindow {
     private static final String TITLE_FORMAT = "SnapshotStartingWindow for taskId=%s";
 
     private static final long DELAY_REMOVAL_TIME_GENERAL = 100;
-    private static final long DELAY_REMOVAL_TIME_IME_VISIBLE = 350;
+    /**
+     * The max delay time in milliseconds for removing the task snapshot window with IME visible.
+     * Ideally the delay time will be shorter when receiving
+     * {@link StartingSurfaceDrawer#onImeDrawnOnTask(int)}.
+     */
+    private static final long MAX_DELAY_REMOVAL_TIME_IME_VISIBLE = 450;
 
     //tmp vars for unused relayout params
     private static final Point TMP_SURFACE_SIZE = new Point();
@@ -138,7 +142,6 @@ public class TaskSnapshotWindow {
     private final RectF mTmpDstFrame = new RectF();
     private final CharSequence mTitle;
     private boolean mHasDrawn;
-    private long mShownTime;
     private boolean mSizeMismatch;
     private final Paint mBackgroundPaint = new Paint();
     private final int mActivityType;
@@ -148,6 +151,8 @@ public class TaskSnapshotWindow {
     private final SurfaceControl.Transaction mTransaction;
     private final Matrix mSnapshotMatrix = new Matrix();
     private final float[] mTmpFloat9 = new float[9];
+    private Runnable mScheduledRunnable;
+    private final boolean mHasImeSurface;
 
     static TaskSnapshotWindow create(StartingWindowInfo info, IBinder appToken,
             TaskSnapshot snapshot, ShellExecutor splashScreenExecutor,
@@ -216,7 +221,7 @@ public class TaskSnapshotWindow {
             taskDescription.setBackgroundColor(WHITE);
         }
 
-        final long delayRemovalTime = snapshot.hasImeSurface() ? DELAY_REMOVAL_TIME_IME_VISIBLE
+        final long delayRemovalTime = snapshot.hasImeSurface() ? MAX_DELAY_REMOVAL_TIME_IME_VISIBLE
                 : DELAY_REMOVAL_TIME_GENERAL;
 
         final TaskSnapshotWindow snapshotSurface = new TaskSnapshotWindow(
@@ -281,12 +286,17 @@ public class TaskSnapshotWindow {
         mDelayRemovalTime = delayRemovalTime;
         mTransaction = new SurfaceControl.Transaction();
         mClearWindowHandler = clearWindowHandler;
+        mHasImeSurface = snapshot.hasImeSurface();
     }
 
     int getBackgroundColor() {
         return mBackgroundPaint.getColor();
     }
 
+    boolean hasImeSurface() {
+        return mHasImeSurface;
+    }
+
     /**
      * Ask system bar background painter to draw status bar background.
      * @hide
@@ -304,21 +314,32 @@ public class TaskSnapshotWindow {
         mSystemBarBackgroundPainter.drawNavigationBarBackground(c);
     }
 
-    void remove() {
-        final long now = SystemClock.uptimeMillis();
-        if ((now - mShownTime < mDelayRemovalTime)
-                // Show the latest content as soon as possible for unlocking to home.
-                && mActivityType != ACTIVITY_TYPE_HOME) {
-            final long delayTime = mShownTime + mDelayRemovalTime - now;
-            mSplashScreenExecutor.executeDelayed(() -> remove(), delayTime);
-            if (DEBUG) {
-                Slog.d(TAG, "Defer removing snapshot surface in " + delayTime);
-            }
+    void scheduleRemove(Runnable onRemove) {
+        // Show the latest content as soon as possible for unlocking to home.
+        if (mActivityType == ACTIVITY_TYPE_HOME) {
+            removeImmediately();
+            onRemove.run();
             return;
         }
+        if (mScheduledRunnable != null) {
+            mSplashScreenExecutor.removeCallbacks(mScheduledRunnable);
+            mScheduledRunnable = null;
+        }
+        mScheduledRunnable = () -> {
+            TaskSnapshotWindow.this.removeImmediately();
+            onRemove.run();
+        };
+        mSplashScreenExecutor.executeDelayed(mScheduledRunnable, mDelayRemovalTime);
+        if (DEBUG) {
+            Slog.d(TAG, "Defer removing snapshot surface in " + mDelayRemovalTime);
+        }
+    }
+
+    void removeImmediately() {
+        mSplashScreenExecutor.removeCallbacks(mScheduledRunnable);
         try {
             if (DEBUG) {
-                Slog.d(TAG, "Removing snapshot surface, mHasDrawn: " + mHasDrawn);
+                Slog.d(TAG, "Removing taskSnapshot surface, mHasDrawn: " + mHasDrawn);
             }
             mSession.remove(mWindow);
         } catch (RemoteException e) {
@@ -356,7 +377,6 @@ public class TaskSnapshotWindow {
         } else {
             drawSizeMatchSnapshot();
         }
-        mShownTime = SystemClock.uptimeMillis();
         mHasDrawn = true;
         reportDrawn();
 
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 d536adb9f8ae..eef0d9bb268f 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
@@ -15,38 +15,53 @@
  */
 package com.android.wm.shell.startingsurface;
 
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
 import static android.window.StartingWindowInfo.STARTING_WINDOW_TYPE_SPLASH_SCREEN;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager;
 import android.content.BroadcastReceiver;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
+import android.graphics.ColorSpace;
+import android.graphics.Point;
 import android.graphics.Rect;
+import android.hardware.HardwareBuffer;
 import android.net.Uri;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.UserHandle;
 import android.testing.TestableContext;
+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.TaskSnapshot;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -61,6 +76,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.MockitoSession;
 
 import java.util.function.IntSupplier;
 
@@ -78,6 +94,7 @@ public class StartingSurfaceDrawerTests {
     private TransactionPool mTransactionPool;
 
     private final Handler mTestHandler = new Handler(Looper.getMainLooper());
+    private ShellExecutor mTestExecutor;
     private final TestableContext mTestContext = new TestContext(
             InstrumentationRegistry.getInstrumentation().getTargetContext());
     TestStartingSurfaceDrawer mStartingSurfaceDrawer;
@@ -138,9 +155,9 @@ public class StartingSurfaceDrawerTests {
 
         doReturn(metrics).when(mMockWindowManager).getMaximumWindowMetrics();
         doNothing().when(mMockWindowManager).addView(any(), any());
-
-        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(mTestContext,
-                new HandlerExecutor(mTestHandler), mTransactionPool));
+        mTestExecutor = new HandlerExecutor(mTestHandler);
+        mStartingSurfaceDrawer = spy(
+                new TestStartingSurfaceDrawer(mTestContext, mTestExecutor, mTransactionPool));
     }
 
     @Test
@@ -205,6 +222,48 @@ public class StartingSurfaceDrawerTests {
         assertEquals(0, windowColor3.mReuseCount);
     }
 
+    @Test
+    public void testRemoveTaskSnapshotWithImeSurfaceWhenOnImeDrawn() throws Exception {
+        final int taskId = 1;
+        final StartingWindowInfo windowInfo =
+                createWindowInfo(taskId, android.R.style.Theme);
+        TaskSnapshot snapshot = createTaskSnapshot(100, 100, new Point(100, 100),
+                new Rect(0, 0, 0, 50), true /* hasImeSurface */);
+        final IWindowSession session = WindowManagerGlobal.getWindowSession();
+        spyOn(session);
+        doReturn(WindowManagerGlobal.ADD_OKAY).when(session).addToDisplay(
+                any() /* window */, any() /* attrs */,
+                anyInt() /* viewVisibility */, anyInt() /* displayId */,
+                any() /* requestedVisibility */, any() /* outInputChannel */,
+                any() /* outInsetsState */, any() /* outActiveControls */);
+        TaskSnapshotWindow mockSnapshotWindow = TaskSnapshotWindow.create(windowInfo,
+                mBinder,
+                snapshot, mTestExecutor, () -> {
+                });
+        spyOn(mockSnapshotWindow);
+        try (AutoCloseable mockTaskSnapshotSession = new AutoCloseable() {
+            MockitoSession mockSession = mockitoSession()
+                    .initMocks(this)
+                    .mockStatic(TaskSnapshotWindow.class)
+                    .startMocking();
+            @Override
+            public void close() {
+                mockSession.finishMocking();
+            }
+        }) {
+            when(TaskSnapshotWindow.create(eq(windowInfo), eq(mBinder), eq(snapshot), any(),
+                    any())).thenReturn(mockSnapshotWindow);
+            // Simulate a task snapshot window created with IME snapshot shown.
+            mStartingSurfaceDrawer.makeTaskSnapshotWindow(windowInfo, mBinder, snapshot);
+            waitHandlerIdle(mTestHandler);
+
+            // Verify the task snapshot with IME snapshot will be removed when received the real IME
+            // drawn callback.
+            mStartingSurfaceDrawer.onImeDrawnOnTask(1);
+            verify(mockSnapshotWindow).removeImmediately();
+        }
+    }
+
     private StartingWindowInfo createWindowInfo(int taskId, int themeResId) {
         StartingWindowInfo windowInfo = new StartingWindowInfo();
         final ActivityInfo info = new ActivityInfo();
@@ -216,10 +275,27 @@ public class StartingSurfaceDrawerTests {
         taskInfo.taskId = taskId;
         windowInfo.targetActivityInfo = info;
         windowInfo.taskInfo = taskInfo;
+        windowInfo.topOpaqueWindowInsetsState = new InsetsState();
+        windowInfo.mainWindowLayoutParams = new WindowManager.LayoutParams();
+        windowInfo.topOpaqueWindowLayoutParams = new WindowManager.LayoutParams();
         return windowInfo;
     }
 
     private static void waitHandlerIdle(Handler handler) {
         handler.runWithScissors(() -> { }, 0 /* timeout */);
     }
+
+    private TaskSnapshot createTaskSnapshot(int width, int height, Point taskSize,
+            Rect contentInsets, boolean hasImeSurface) {
+        final HardwareBuffer buffer = HardwareBuffer.create(width, height, HardwareBuffer.RGBA_8888,
+                1, HardwareBuffer.USAGE_CPU_READ_RARELY);
+        return new TaskSnapshot(
+                System.currentTimeMillis(),
+                new ComponentName("", ""), buffer,
+                ColorSpace.get(ColorSpace.Named.SRGB), ORIENTATION_PORTRAIT,
+                Surface.ROTATION_0, taskSize, contentInsets, false,
+                true /* isRealSnapshot */, WINDOWING_MODE_FULLSCREEN,
+                0 /* systemUiVisibility */, false /* isTranslucent */,
+                hasImeSurface /* hasImeSurface */);
+    }
 }
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index c2cfe0b7bd92..6e6cdf24bb1b 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -2253,6 +2253,17 @@ 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() {
         removeStartingWindowAnimation(true /* prepareAnimation */);
     }
@@ -5824,7 +5835,7 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
             // own stuff.
             win.cancelAnimation();
         }
-        removeStartingWindow();
+        removeStartingWindowIfNeeded();
         updateReportedVisibilityLocked();
     }
 
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index da47328691c0..ed1e784bf275 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -90,6 +90,24 @@ final class ImeInsetsSourceProvider extends InsetsSourceProvider {
         onSourceChanged();
     }
 
+    @Override
+    protected boolean updateClientVisibility(InsetsControlTarget caller) {
+        boolean changed = super.updateClientVisibility(caller);
+        if (changed && caller.getRequestedVisibility(mSource.getType())) {
+            reportImeDrawnForOrganizer(caller);
+        }
+        return changed;
+    }
+
+    private void reportImeDrawnForOrganizer(InsetsControlTarget caller) {
+        if (caller.getWindow() != null && caller.getWindow().getTask() != null) {
+            if (caller.getWindow().getTask().isOrganized()) {
+                mWin.mWmService.mAtmService.mTaskOrganizerController.reportImeDrawnOnTask(
+                        caller.getWindow().getTask());
+            }
+        }
+    }
+
     private void onSourceChanged() {
         if (mLastSource.equals(mSource)) {
             return;
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index ced5af126e49..364594e4286a 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -5144,7 +5144,7 @@ class Task extends WindowContainer<WindowContainer> {
     /**
      * @return true if the task is currently focused.
      */
-    private boolean isFocused() {
+    boolean isFocused() {
         if (mDisplayContent == null || mDisplayContent.mCurrentFocus == null) {
             return false;
         }
@@ -5206,6 +5206,10 @@ class Task extends WindowContainer<WindowContainer> {
      * @param hasFocus
      */
     void onWindowFocusChanged(boolean hasFocus) {
+        final ActivityRecord topAct = getTopVisibleActivity();
+        if (topAct != null && (topAct.mStartingData instanceof SnapshotStartingData)) {
+            topAct.removeStartingWindowIfNeeded();
+        }
         updateShadowsRadius(hasFocus, getSyncTransaction());
         // TODO(b/180525887): Un-comment once there is resolution on the bug.
         // 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 09c5581385dd..88467baa6c34 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -733,6 +733,17 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
         mPendingTaskEvents.clear();
     }
 
+    void reportImeDrawnOnTask(Task task) {
+        final TaskOrganizerState state = mTaskOrganizerStates.get(task.mTaskOrganizer.asBinder());
+        if (state != null) {
+            try {
+                state.mOrganizer.mTaskOrganizer.onImeDrawnOnTask(task.mTaskId);
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Exception sending onImeDrawnOnTask callback", e);
+            }
+        }
+    }
+
     void onTaskInfoChanged(Task task, boolean force) {
         if (!task.mTaskAppearedSent) {
             // Skip if task still not appeared.
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 d6a8401f5b18..ab496cf34acc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -797,6 +797,9 @@ public class WindowOrganizerTests extends WindowTestsBase {
         public void onBackPressedOnTaskRoot(RunningTaskInfo taskInfo) {
         }
         @Override
+        public void onImeDrawnOnTask(int taskId) throws RemoteException {
+        }
+        @Override
         public void onAppSplashScreenViewRemoved(int taskId) {
         }
     };