diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 2cebc3be4b71e..79346e057d9d0 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -647,6 +647,9 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
     private boolean mLastContainsDismissKeyguardWindow;
     private boolean mLastContainsTurnScreenOnWindow;
 
+    /** Whether the IME is showing when transitioning away from this activity. */
+    boolean mLastImeShown;
+
     /**
      * A flag to determine if this AR is in the process of closing or entering PIP. This is needed
      * to help AR know that the app is in the process of closing but hasn't yet started closing on
@@ -4722,6 +4725,15 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
             setClientVisible(visible);
         }
 
+        if (!visible) {
+            final InsetsControlTarget imeInputTarget = mDisplayContent.getImeTarget(
+                    DisplayContent.IME_TARGET_INPUT);
+            mLastImeShown = imeInputTarget != null && imeInputTarget.getWindow() != null
+                    && imeInputTarget.getWindow().mActivityRecord == this
+                    && mDisplayContent.mInputMethodWindow != null
+                    && mDisplayContent.mInputMethodWindow.isVisible();
+        }
+
         final DisplayContent displayContent = getDisplayContent();
         if (!displayContent.mClosingApps.contains(this)
                 && !displayContent.mOpeningApps.contains(this)) {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index e0bae9d8de32c..23dff87c4418f 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -63,6 +63,8 @@ import static android.view.WindowManager.LayoutParams.FLAG_NOT_TOUCHABLE;
 import static android.view.WindowManager.LayoutParams.FLAG_NOT_TOUCH_MODAL;
 import static android.view.WindowManager.LayoutParams.FLAG_SPLIT_TOUCH;
 import static android.view.WindowManager.LayoutParams.LAST_APPLICATION_WINDOW;
+import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_HIDDEN;
+import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_HIDDEN;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_BOOT_PROGRESS;
@@ -208,6 +210,7 @@ import android.view.Surface.Rotation;
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Transaction;
 import android.view.SurfaceSession;
+import android.view.View;
 import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.view.WindowManager.DisplayImePolicy;
@@ -1549,9 +1552,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             // to cover the activity configuration change.
             return false;
         }
-        if ((r.mStartingData != null && r.mStartingData.hasImeSurface())
-                || (mInsetsStateController.getImeSourceProvider()
-                        .getSource().getVisibleFrame() != null)) {
+        if (r.attachedToProcess() && mayImeShowOnLaunchingActivity(r)) {
             // Currently it is unknown that when will IME window be ready. Reject the case to
             // avoid flickering by showing IME in inconsistent orientation.
             return false;
@@ -1607,6 +1608,24 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         return true;
     }
 
+    /** Returns {@code true} if the IME is possible to show on the launching activity. */
+    private boolean mayImeShowOnLaunchingActivity(@NonNull ActivityRecord r) {
+        final WindowState win = r.findMainWindow();
+        if (win == null) {
+            return false;
+        }
+        // See InputMethodManagerService#shouldRestoreImeVisibility that we expecting the IME
+        // should be hidden when the window set the hidden softInputMode.
+        final int softInputMode = win.mAttrs.softInputMode;
+        switch (softInputMode & WindowManager.LayoutParams.SOFT_INPUT_MASK_STATE) {
+            case SOFT_INPUT_STATE_ALWAYS_HIDDEN:
+            case SOFT_INPUT_STATE_HIDDEN:
+                return false;
+        }
+        return r.mLastImeShown && mInputMethodWindow != null && mInputMethodWindow.mHasSurface
+                && mInputMethodWindow.mViewVisibility == View.VISIBLE;
+    }
+
     /** Returns {@code true} if the top activity is transformed with the new rotation of display. */
     boolean hasTopFixedRotationLaunchingApp() {
         return mFixedRotationLaunchingApp != null
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 9ffb2b1769153..0000f95e7849d 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -443,7 +443,8 @@ class TaskSnapshotController {
         } else {
             excludeLayers = new SurfaceControl[0];
         }
-        builder.setHasImeSurface(!excludeIme && imeWindow != null && imeWindow.isDrawn());
+        builder.setHasImeSurface(!excludeIme && imeWindow != null && imeWindow.isVisible());
+
         final SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
                 SurfaceControl.captureLayersExcluding(
                         task.getSurfaceControl(), mTmpRect, scaleFraction,
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index e9e201327aa46..97afc165624ac 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -214,8 +214,8 @@ public class TaskSnapshotControllerTest extends WindowTestsBase {
         spyOn(mDisplayContent);
         spyOn(mDisplayContent.mInputMethodWindow);
         when(task.getDisplayContent().shouldImeAttachedToApp()).thenReturn(true);
-        // Intentionally set the IME window is in drawn state.
-        doReturn(true).when(mDisplayContent.mInputMethodWindow).isDrawn();
+        // Intentionally set the IME window is in visible state.
+        doReturn(true).when(mDisplayContent.mInputMethodWindow).isVisible();
         // Verify no NPE happens when calling createTaskSnapshot.
         try {
             final TaskSnapshot.Builder builder = new TaskSnapshot.Builder();