diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index da47328691c0a..4f6a693b8c3fc 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -29,6 +29,7 @@ import static com.android.server.wm.ImeInsetsSourceProviderProto.IS_IME_LAYOUT_D
 import static com.android.server.wm.WindowManagerService.H.UPDATE_MULTI_WINDOW_STACKS;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.os.Trace;
 import android.util.proto.ProtoOutputStream;
 import android.view.InsetsSource;
@@ -90,6 +91,16 @@ final class ImeInsetsSourceProvider extends InsetsSourceProvider {
         onSourceChanged();
     }
 
+    @Override
+    void updateControlForTarget(@Nullable InsetsControlTarget target, boolean force) {
+        if (target != null && target.getWindow() != null) {
+            // ime control target could be a different window.
+            // Refer WindowState#getImeControlTarget().
+            target = target.getWindow().getImeControlTarget();
+        }
+        super.updateControlForTarget(target, force);
+    }
+
     private void onSourceChanged() {
         if (mLastSource.equals(mSource)) {
             return;
diff --git a/services/core/java/com/android/server/wm/InsetsSourceProvider.java b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
index 7daebff2ccc20..cbd1314b104a5 100644
--- a/services/core/java/com/android/server/wm/InsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
@@ -307,11 +307,6 @@ class InsetsSourceProvider {
             // to control the window for now.
             return;
         }
-        if (target != null && target.getWindow() != null) {
-            // ime control target could be a different window.
-            // Refer WindowState#getImeControlTarget().
-            target = target.getWindow().getImeControlTarget();
-        }
 
         if (mWin != null && mWin.getSurfaceControl() == null) {
             // if window doesn't have a surface, set it null and return.