diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index de6501aa9718d..5d3f3c0401ef0 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -3538,9 +3538,9 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
     InsetsControlTarget getImeFallback() {
         // host is in non-default display that doesn't support system decor, default to
         // default display's StatusBar to control IME (when available), else let system control it.
-        WindowState statusBar = 
-                mWmService.getDefaultDisplayContentLocked().getDisplayPolicy().getStatusBar();
-        return statusBar != null ? statusBar : mRemoteInsetsControlTarget;
+        final DisplayContent defaultDc = mWmService.getDefaultDisplayContentLocked();
+        WindowState statusBar = defaultDc.getDisplayPolicy().getStatusBar();
+        return statusBar != null ? statusBar : defaultDc.mRemoteInsetsControlTarget;
     }
 
     boolean canShowIme() {
@@ -3600,7 +3600,10 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
      */
     @VisibleForTesting
     InsetsControlTarget computeImeControlTarget() {
-        if (!isImeControlledByApp() && mRemoteInsetsControlTarget != null) {
+        if (!isImeControlledByApp() && mRemoteInsetsControlTarget != null
+                || (mInputMethodInputTarget != null
+                        && getImeHostOrFallback(mInputMethodInputTarget.getWindow())
+                                == mRemoteInsetsControlTarget)) {
             return mRemoteInsetsControlTarget;
         } else {
             // Now, a special case -- if the last target's window is in the process of exiting, but
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index 99ee5e121b7a4..e7fbc334306e7 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -32,7 +32,7 @@ import java.io.PrintWriter;
  */
 class ImeInsetsSourceProvider extends InsetsSourceProvider {
 
-    private WindowState mImeTargetFromIme;
+    private InsetsControlTarget mImeTargetFromIme;
     private Runnable mShowImeRunner;
     private boolean mIsImeLayoutDrawn;
 
@@ -47,10 +47,12 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
      *
      * @param imeTarget imeTarget on which IME request is coming from.
      */
-    void scheduleShowImePostLayout(WindowState imeTarget) {
+    void scheduleShowImePostLayout(InsetsControlTarget imeTarget) {
         boolean targetChanged = mImeTargetFromIme != imeTarget
                 && mImeTargetFromIme != null && imeTarget != null && mShowImeRunner != null
-                && mImeTargetFromIme.mActivityRecord == imeTarget.mActivityRecord;
+                && imeTarget.getWindow() != null && mImeTargetFromIme.getWindow() != null
+                && mImeTargetFromIme.getWindow().mActivityRecord
+                        == imeTarget.getWindow().mActivityRecord;
         mImeTargetFromIme = imeTarget;
         if (targetChanged) {
             // target changed, check if new target can show IME.
@@ -62,7 +64,8 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
             return;
         }
 
-        ProtoLog.d(WM_DEBUG_IME, "Schedule IME show for %s", mImeTargetFromIme.getName());
+        ProtoLog.d(WM_DEBUG_IME, "Schedule IME show for %s", mImeTargetFromIme.getWindow() == null
+                ? mImeTargetFromIme : mImeTargetFromIme.getWindow().getName());
         mShowImeRunner = () -> {
             ProtoLog.d(WM_DEBUG_IME, "Run showImeRunner");
             // Target should still be the same.
@@ -127,13 +130,17 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
             return false;
         }
         ProtoLog.d(WM_DEBUG_IME, "dcTarget: %s mImeTargetFromIme: %s",
-                dcTarget.getName(), mImeTargetFromIme.getName());
+                dcTarget.getName(), mImeTargetFromIme.getWindow() == null
+                        ? mImeTargetFromIme : mImeTargetFromIme.getWindow().getName());
 
         return (!dcTarget.isClosing() && mImeTargetFromIme == dcTarget)
-                || (mImeTargetFromIme != null && dcTarget.getParentWindow() == mImeTargetFromIme
-                        && dcTarget.mSubLayer > mImeTargetFromIme.mSubLayer)
+                || (mImeTargetFromIme != null && mImeTargetFromIme.getWindow() != null
+                        && dcTarget.getParentWindow() == mImeTargetFromIme
+                        && dcTarget.mSubLayer > mImeTargetFromIme.getWindow().mSubLayer)
                 || mImeTargetFromIme == mDisplayContent.getImeFallback()
-                || (!mImeTargetFromIme.isClosing() && controlTarget == mImeTargetFromIme);
+                || controlTarget == mImeTargetFromIme
+                        && (mImeTargetFromIme.getWindow() == null 
+                                || !mImeTargetFromIme.getWindow().isClosing());
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 0b50c1cd496a7..ef81c0a5d206c 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7606,13 +7606,14 @@ public class WindowManagerService extends IWindowManager.Stub
                 if (imeTarget == null) {
                     return;
                 }
-                imeTarget = imeTarget.getImeControlTarget().getWindow();
+                final InsetsControlTarget controlTarget = imeTarget.getImeControlTarget();
+                imeTarget = controlTarget.getWindow();
                 // If InsetsControlTarget doesn't have a window, its using remoteControlTarget which
                 // is controlled by default display
                 final DisplayContent dc = imeTarget != null
                         ? imeTarget.getDisplayContent() : getDefaultDisplayContentLocked();
                 dc.getInsetsStateController().getImeSourceProvider()
-                        .scheduleShowImePostLayout(imeTarget);
+                        .scheduleShowImePostLayout(controlTarget);
             }
         }