diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index f71d08afb1cb0..3b6c733776728 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2297,6 +2297,27 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
         return mLastConfigReportedToClient;
     }
 
+    @Override
+    public void onConfigurationChanged(Configuration newParentConfig) {
+        if (getDisplayContent().getImeTarget(IME_TARGET_INPUT) != this && !isImeLayeringTarget()) {
+            super.onConfigurationChanged(newParentConfig);
+            return;
+        }
+
+        mTempConfiguration.setTo(getConfiguration());
+        super.onConfigurationChanged(newParentConfig);
+        final boolean windowConfigChanged = mTempConfiguration.windowConfiguration
+                .diff(newParentConfig.windowConfiguration, false) != 0;
+
+        // When the window configuration changed, we need to update the IME control target in
+        // case the app may lose the IME inets control when exiting from split-screen mode, or the
+        // IME parent may failed to attach to the app during rotating the screen.
+        // See DisplayContent#isImeAttachedToApp, DisplayContent#isImeControlledByApp
+        if (windowConfigChanged) {
+            getDisplayContent().updateImeControlTarget();
+        }
+    }
+
     @Override
     void onMergedOverrideConfigurationChanged() {
         super.onMergedOverrideConfigurationChanged();
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
index bfbe203fb65e0..8d067beff44cc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
@@ -17,6 +17,7 @@
 package com.android.server.wm;
 
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_SPLIT_SCREEN_PRIMARY;
 import static android.view.InsetsState.ITYPE_STATUS_BAR;
@@ -47,6 +48,7 @@ import static com.android.dx.mockito.inline.extended.ExtendedMockito.reset;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static com.android.server.wm.DisplayContent.IME_TARGET_CONTROL;
 import static com.android.server.wm.WindowContainer.SYNC_STATE_WAITING_FOR_DRAW;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -877,4 +879,30 @@ public class WindowStateTests extends WindowTestsBase {
         mDisplayContent.getInsetsStateController().notifyInsetsChanged();
         verify(app).notifyInsetsChanged();
     }
+
+    @UseTestDisplay(addWindows = { W_ACTIVITY })
+    @Test
+    public void testUpdateImeControlTargetWhenLeavingMultiWindow() {
+        WindowState app = createWindow(null, TYPE_BASE_APPLICATION,
+                mAppWindow.mToken, "app");
+        mDisplayContent.setRemoteInsetsController(createDisplayWindowInsetsController());
+
+        spyOn(app);
+        mDisplayContent.setImeInputTarget(mAppWindow);
+        mDisplayContent.setImeLayeringTarget(mAppWindow);
+
+        // Simulate entering multi-window mode and verify if the IME control target is remote.
+        app.mActivityRecord.getRootTask().setWindowingMode(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY);
+        assertEquals(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY, app.getWindowingMode());
+        assertEquals(mDisplayContent.mRemoteInsetsControlTarget,
+                mDisplayContent.computeImeControlTarget());
+
+        // Simulate exiting multi-window mode and verify if the IME control target changed
+        // to the app window.
+        spyOn(app.getDisplayContent());
+        app.mActivityRecord.getRootTask().setWindowingMode(WINDOWING_MODE_FULLSCREEN);
+
+        verify(app.getDisplayContent()).updateImeControlTarget();
+        assertEquals(mAppWindow, mDisplayContent.getImeTarget(IME_TARGET_CONTROL).getWindow());
+    }
 }