diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 2c315445577d3..2f9868e059d3b 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -1588,9 +1588,24 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         // TODO(multidisplay): Needs some serious rethought when the target and IME are not on the
         // same display. Or even when the current IME/target are not on the same screen as the next
         // IME/target. For now only look for input windows on the main screen.
-        mUpdateImeTarget = updateImeTarget;
-        WindowState target = getWindow(mComputeImeTargetPredicate);
 
+        // The target candidate provided by the IME tells us which window token, but not which
+        // window within the token (e.g. child windows...). So, we use the token to look-up the
+        // best target window.
+        // TODO: Have the input method service report the right window with the token vs. just the
+        // base window of the token.
+        final WindowState baseWin = mService.getWindow(mService.mInputMethodTargetCandidate);
+        final WindowToken targetToken = baseWin != null ? baseWin.mToken : null;
+        WindowState target = targetToken != null ?
+                targetToken.getWindow(mComputeImeTargetPredicate) : null;
+        // If there isn't a better candidate in the token (maybe because they are not visible), then
+        // fall back to targeting the base window of the token, so the IME can still maintain the
+        // right z-order based on the last person that set it vs. changing its z-order to the very
+        // up since there if target is null.
+        // TODO: Consider z-ordering IME to bottom instead of top if there is no visible target.
+        // Also, consider tying the visible the visibility of the IME to the current target. I.e if
+        // target isn't visible, then IME shouldn't be visible.
+        target = target == null ? baseWin : target;
 
         // Yet more tricksyness!  If this window is a "starting" window, we do actually want
         // to be on top of it, but it is not -really- where input will go. So look down below
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 474610bc6331f..5795cb854ed70 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -16,7 +16,6 @@
 
 package com.android.server.wm;
 
-import static android.Manifest.permission.INTERNAL_SYSTEM_WINDOW;
 import static android.Manifest.permission.MANAGE_APP_TOKENS;
 import static android.Manifest.permission.REGISTER_WINDOW_MANAGER_LISTENERS;
 import static android.app.ActivityManager.DOCKED_STACK_CREATE_MODE_TOP_OR_LEFT;
@@ -30,7 +29,6 @@ import static android.content.Intent.EXTRA_PACKAGE_NAME;
 import static android.content.Intent.EXTRA_UID;
 import static android.content.Intent.EXTRA_USER_HANDLE;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
-import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.os.UserHandle.USER_NULL;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.DOCKED_INVALID;
@@ -69,8 +67,6 @@ import static android.view.WindowManagerPolicy.FINISH_LAYOUT_REDO_WALLPAPER;
 import static com.android.server.wm.AppTransition.TRANSIT_UNSET;
 import static com.android.server.wm.AppWindowAnimator.PROLONG_ANIMATION_AT_END;
 import static com.android.server.wm.AppWindowAnimator.PROLONG_ANIMATION_AT_START;
-import static com.android.server.wm.DragResizeMode.DRAG_RESIZE_MODE_DOCKED_DIVIDER;
-import static com.android.server.wm.DragResizeMode.DRAG_RESIZE_MODE_FREEFORM;
 import static com.android.server.wm.KeyguardDisableHandler.KEYGUARD_POLICY_CHANGED;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_ADD_REMOVE;
@@ -112,7 +108,6 @@ import android.app.ActivityManager.TaskSnapshot;
 import android.app.ActivityManagerInternal;
 import android.app.AppOpsManager;
 import android.app.IActivityManager;
-import android.app.admin.DevicePolicyManager;
 import android.content.BroadcastReceiver;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -626,9 +621,14 @@ public class WindowManagerService extends IWindowManager.Stub
     WindowState mCurrentFocus = null;
     WindowState mLastFocus = null;
 
-    /** This just indicates the window the input method is on top of, not
-     * necessarily the window its input is going to. */
+    // TODO: All the IME window tracking should be moved to DisplayContent and tracked per display.
+    // This just indicates the window the input method is on top of, not necessarily the window its
+    // input is going to.
     WindowState mInputMethodTarget = null;
+    // The binder token currently using the IME as determined by the input method service.
+    // Window manager uses this to determine the final input method target
+    // (almost always this candidate) for z-ordering.
+    IBinder mInputMethodTargetCandidate = null;
 
     /** If true hold off on modifying the animation layer of mInputMethodTarget */
     boolean mInputMethodTargetWaitingAnim;
@@ -6951,8 +6951,9 @@ public class WindowManagerService extends IWindowManager.Stub
             pw.print("  mLastFocus="); pw.println(mLastFocus);
         }
         pw.print("  mFocusedApp="); pw.println(mFocusedApp);
-        if (mInputMethodTarget != null) {
-            pw.print("  mInputMethodTarget="); pw.println(mInputMethodTarget);
+        if (mInputMethodTarget != null || mInputMethodTargetCandidate != null) {
+            pw.println("  mInputMethodTarget=" + mInputMethodTarget
+                    + " mInputMethodTargetCandidate=" + getWindow(mInputMethodTargetCandidate));
         }
         pw.print("  mInTouchMode="); pw.print(mInTouchMode);
                 pw.print(" mLayoutSeq="); pw.println(mLayoutSeq);
@@ -7819,11 +7820,28 @@ public class WindowManagerService extends IWindowManager.Stub
         @Override
         public void updateInputMethodWindowStatus(@NonNull IBinder imeToken,
                 boolean imeWindowVisible, @Nullable IBinder targetWindowToken) {
-            // TODO (b/34628091): Use this method to address the window animation issue.
-            if (DEBUG_INPUT_METHOD) {
-                Slog.w(TAG_WM, "updateInputMethodWindowStatus: imeToken=" + imeToken
-                        + " imeWindowVisible=" + imeWindowVisible
-                        + " targetWindowToken=" + targetWindowToken);
+            synchronized (mWindowMap) {
+                final WindowState newTargetWin = getWindow(targetWindowToken);
+                final WindowState currentTargetWin = getWindow(mInputMethodTargetCandidate);
+
+                if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "updateInputMethodWindowStatus: imeToken="
+                        + imeToken + " imeWindowVisible=" + imeWindowVisible
+                        + " targetWindowToken=" + targetWindowToken
+                        + " newTargetWin=" + newTargetWin
+                        + " currentTargetWin=" + currentTargetWin);
+
+                if (newTargetWin == currentTargetWin) {
+                    return;
+                }
+
+                final DisplayContent dc = newTargetWin != null
+                        ? newTargetWin.getDisplayContent() : currentTargetWin.getDisplayContent();
+
+                // It is possible the window for the target candidate isn't added yet, so we
+                // remember the token instead and use it to look-up the window each time we compute
+                // the ime target.
+                mInputMethodTargetCandidate = targetWindowToken;
+                dc.computeImeTarget(true /* updateImeTarget */);
             }
         }
 
@@ -7869,6 +7887,10 @@ public class WindowManagerService extends IWindowManager.Stub
         }
     }
 
+    WindowState getWindow(IBinder binder) {
+        return binder == null ? null : mWindowMap.get(binder);
+    }
+
     void registerAppFreezeListener(AppFreezeListener listener) {
         if (!mAppFreezeListeners.contains(listener)) {
             mAppFreezeListeners.add(listener);
diff --git a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
index bd3271b327158..73ad7c241396e 100644
--- a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
@@ -22,6 +22,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_ATTACHED_
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_VOICE_INTERACTION;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 
 import org.junit.Test;
@@ -146,6 +147,7 @@ public class DisplayContentTests extends WindowTestsBase {
         final WindowState appWin = createWindow(null, TYPE_APPLICATION, sDisplayContent, "appWin");
         appWin.setHasSurface(true);
         assertTrue(appWin.canBeImeTarget());
+        sWm.mInputMethodTargetCandidate = appWin.mClient.asBinder();
         WindowState imeTarget = sDisplayContent.computeImeTarget(false /* updateImeTarget */);
         assertEquals(appWin, imeTarget);
 
@@ -156,6 +158,20 @@ public class DisplayContentTests extends WindowTestsBase {
         assertTrue(childWin.canBeImeTarget());
         imeTarget = sDisplayContent.computeImeTarget(false /* updateImeTarget */);
         assertEquals(childWin, imeTarget);
+
+        final WindowState appWin2 =
+                createWindow(null, TYPE_APPLICATION, sDisplayContent, "appWin2");
+        appWin2.setHasSurface(true);
+        assertTrue(appWin2.canBeImeTarget());
+        // Verify that the IME target isn't adjusted since mInputMethodTargetCandidate didn't change
+        // to the new app.
+        imeTarget = sDisplayContent.computeImeTarget(false /* updateImeTarget */);
+        assertNotEquals(appWin2, imeTarget);
+
+        sWm.mInputMethodTargetCandidate = appWin2.mClient.asBinder();
+        // Verify app is not IME target since its token is set as a candidate.
+        imeTarget = sDisplayContent.computeImeTarget(false /* updateImeTarget */);
+        assertEquals(appWin2, imeTarget);
     }
 
     /**
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
index 52e10a542238f..b9c2eedfa2158 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -69,7 +69,6 @@ import java.util.LinkedList;
 class WindowTestsBase {
     static WindowManagerService sWm = null;
     static TestWindowManagerPolicy sPolicy = null;
-    private final static IWindow sIWindow = new TestIWindow();
     private final static Session sMockSession = mock(Session.class);
     private static int sNextDisplayId = Display.DEFAULT_DISPLAY + 1;
     static int sNextStackId = FIRST_DYNAMIC_STACK_ID;
@@ -148,6 +147,7 @@ class WindowTestsBase {
         }
 
         sWm.mInputMethodTarget = null;
+        sWm.mInputMethodTargetCandidate = null;
     }
 
     private static WindowState createCommonWindow(WindowState parent, int type, String name) {
@@ -237,11 +237,12 @@ class WindowTestsBase {
         final WindowManager.LayoutParams attrs = new WindowManager.LayoutParams(type);
         attrs.setTitle(name);
 
-        final WindowState w = new WindowState(sWm, sMockSession, sIWindow, token, parent, OP_NONE,
-                0, attrs, 0, 0, ownerCanAddInternalSystemWindow);
+        final WindowState w = new WindowState(sWm, sMockSession, new TestIWindow(), token, parent,
+                OP_NONE, 0, attrs, 0, 0, ownerCanAddInternalSystemWindow);
         // TODO: Probably better to make this call in the WindowState ctor to avoid errors with
         // adding it to the token...
         token.addWindow(w);
+        sWm.mWindowMap.put(w.mClient.asBinder(), w);
         return w;
     }
 
@@ -463,8 +464,9 @@ class WindowTestsBase {
         boolean resizeReported;
 
         TestWindowState(WindowManager.LayoutParams attrs, WindowToken token) {
-            super(sWm, sMockSession, sIWindow, token, null, OP_NONE, 0, attrs, 0, 0,
+            super(sWm, sMockSession, new TestIWindow(), token, null, OP_NONE, 0, attrs, 0, 0,
                     false /* ownerCanAddInternalSystemWindow */);
+            sWm.mWindowMap.put(mClient.asBinder(), this);
         }
 
         @Override