diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index 5647bf9..c5a11abe 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -605,9 +605,6 @@
             if (DEBUG) Log.v(TAG, "unbindInput(): binding=" + mInputBinding
                     + " ic=" + mInputConnection);
             // Unbind input is per process per display.
-            // TODO(b/150902448): free-up IME surface when target is changing.
-            //  e.g. DisplayContent#setInputMethodTarget()
-            removeImeSurface();
             onUnbindInput();
             mInputBinding = null;
             mInputConnection = null;
diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index ef9d990..c1998c6 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -119,11 +119,11 @@
         // If we had a request before to show from IME (tracked with mImeRequestedShow), reaching
         // this code here means that we now got control, so we can start the animation immediately.
         // If client window is trying to control IME and IME is already visible, it is immediate.
-        if (fromIme || mState.getSource(getType()).isVisible()) {
+        if (fromIme || mState.getSource(getType()).isVisible() && getControl() != null) {
             return ShowResult.SHOW_IMMEDIATELY;
         }
 
-        return getImm().requestImeShow(null /* resultReceiver */)
+        return getImm().requestImeShow(mController.getHost().getWindowToken())
                 ? ShowResult.IME_SHOW_DELAYED : ShowResult.IME_SHOW_FAILED;
     }
 
@@ -132,12 +132,15 @@
      */
     @Override
     void notifyHidden() {
-        getImm().notifyImeHidden();
+        getImm().notifyImeHidden(mController.getHost().getWindowToken());
     }
 
     @Override
     public void removeSurface() {
-        getImm().removeImeSurface();
+        final IBinder window = mController.getHost().getWindowToken();
+        if (window != null) {
+            getImm().removeImeSurface(window);
+        }
     }
 
     @Override
@@ -146,6 +149,7 @@
         super.setControl(control, showTypes, hideTypes);
         if (control == null && !mIsRequestedVisibleAwaitingControl) {
             hide();
+            removeSurface();
         }
     }
 
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index 40e6f4b..700dc66 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -153,6 +153,9 @@
                 if (oldLeash == null || newLeash == null || !oldLeash.isSameSurface(newLeash)) {
                     applyHiddenToControl();
                 }
+                if (!requestedVisible && !mIsAnimationPending) {
+                    removeSurface();
+                }
             }
         }
         if (lastControl != null) {
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index a586425..2b7044d 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -2725,7 +2725,6 @@
                             mAttachInfo.mThreadedRenderer.isEnabled()) {
                         mAttachInfo.mThreadedRenderer.destroy();
                     }
-                    notifySurfaceDestroyed();
                 } else if ((surfaceReplaced
                         || surfaceSizeChanged || windowRelayoutWasForced || colorModeChanged)
                         && mSurfaceHolder == null
@@ -2956,6 +2955,10 @@
             }
         }
 
+        if (surfaceDestroyed) {
+            notifySurfaceDestroyed();
+        }
+
         if (triggerGlobalLayoutListener) {
             mAttachInfo.mRecomputeGlobalAttributes = false;
             mAttachInfo.mTreeObserver.dispatchOnGlobalLayout();
@@ -9377,6 +9380,11 @@
         return mInputEventReceiver.getToken();
     }
 
+    @NonNull
+    public IBinder getWindowToken() {
+        return mAttachInfo.mWindowToken;
+    }
+
     /**
      * Class for managing the accessibility interaction connection
      * based on the global accessibility state.
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index aedb59b..3be0a4d 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -2109,28 +2109,36 @@
 
     /**
      * Call showSoftInput with currently focused view.
-     * @return {@code true} if IME can be shown.
+     *
+     * @param windowToken the window from which this request originates. If this doesn't match the
+     *                    currently served view, the request is ignored and returns {@code false}.
+     *
+     * @return {@code true} if IME can (eventually) be shown, {@code false} otherwise.
      * @hide
      */
-    public boolean requestImeShow(ResultReceiver resultReceiver) {
+    public boolean requestImeShow(IBinder windowToken) {
         synchronized (mH) {
             final View servedView = getServedViewLocked();
-            if (servedView == null) {
+            if (servedView == null || servedView.getWindowToken() != windowToken) {
                 return false;
             }
-            showSoftInput(servedView, 0 /* flags */, resultReceiver);
+            showSoftInput(servedView, 0 /* flags */, null /* resultReceiver */);
             return true;
         }
     }
 
     /**
      * Notify IME directly that it is no longer visible.
+     *
+     * @param windowToken the window from which this request originates. If this doesn't match the
+     *                    currently served view, the request is ignored.
      * @hide
      */
-    public void notifyImeHidden() {
+    public void notifyImeHidden(IBinder windowToken) {
         synchronized (mH) {
             try {
-                if (mCurMethod != null) {
+                if (mCurMethod != null && mCurRootView != null
+                        && mCurRootView.getWindowToken() == windowToken) {
                     mCurMethod.notifyImeHidden();
                 }
             } catch (RemoteException re) {
@@ -2140,15 +2148,15 @@
 
     /**
      * Notify IME directly to remove surface as it is no longer visible.
+     * @param windowToken The client window token that requests the IME to remove its surface.
      * @hide
      */
-    public void removeImeSurface() {
+    public void removeImeSurface(IBinder windowToken) {
         synchronized (mH) {
             try {
-                if (mCurMethod != null) {
-                    mCurMethod.removeImeSurface();
-                }
-            } catch (RemoteException re) {
+                mService.removeImeSurfaceFromWindow(windowToken);
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
             }
         }
     }
diff --git a/core/java/com/android/internal/view/IInputMethodManager.aidl b/core/java/com/android/internal/view/IInputMethodManager.aidl
index 8ec51b8..a1cbd3f 100644
--- a/core/java/com/android/internal/view/IInputMethodManager.aidl
+++ b/core/java/com/android/internal/view/IInputMethodManager.aidl
@@ -73,5 +73,8 @@
             in float[] matrixValues);
 
     oneway void reportPerceptible(in IBinder windowToken, boolean perceptible);
+    /** Remove the IME surface. Requires INTERNAL_SYSTEM_WINDOW permission. */
     void removeImeSurface();
+    /** Remove the IME surface. Requires passing the currently focused window. */
+    void removeImeSurfaceFromWindow(in IBinder windowToken);
 }
diff --git a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
index 1b32725..7efd616 100644
--- a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
@@ -71,6 +71,9 @@
     private SurfaceControl mLeash;
     @Mock Transaction mMockTransaction;
     private InsetsSource mSpyInsetsSource;
+    private boolean mRemoveSurfaceCalled = false;
+    private InsetsController mController;
+    private InsetsState mState;
 
     @Before
     public void setup() {
@@ -89,13 +92,19 @@
             } catch (BadTokenException e) {
                 // activity isn't running, lets ignore BadTokenException.
             }
-            InsetsState state = new InsetsState();
+            mState = new InsetsState();
             mSpyInsetsSource = Mockito.spy(new InsetsSource(ITYPE_STATUS_BAR));
-            state.addSource(mSpyInsetsSource);
+            mState.addSource(mSpyInsetsSource);
 
-            mConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, state,
-                    () -> mMockTransaction,
-                    new InsetsController(new ViewRootInsetsControllerHost(viewRootImpl)));
+            mController = new InsetsController(new ViewRootInsetsControllerHost(viewRootImpl));
+            mConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, mState,
+                    () -> mMockTransaction, mController) {
+                @Override
+                public void removeSurface() {
+                    super.removeSurface();
+                    mRemoveSurfaceCalled = true;
+                }
+            };
         });
         instrumentation.waitForIdleSync();
 
@@ -171,6 +180,25 @@
             mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()),
                     new int[1], hideTypes);
             assertEquals(statusBars(), hideTypes[0]);
+            assertFalse(mRemoveSurfaceCalled);
         });
     }
+
+    @Test
+    public void testRestore_noAnimation() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            mConsumer.hide();
+            mController.onStateChanged(mState);
+            mConsumer.setControl(null, new int[1], new int[1]);
+            reset(mMockTransaction);
+            verifyZeroInteractions(mMockTransaction);
+            mRemoveSurfaceCalled = false;
+            int[] hideTypes = new int[1];
+            mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()),
+                    new int[1], hideTypes);
+            assertTrue(mRemoveSurfaceCalled);
+            assertEquals(0, hideTypes[0]);
+        });
+
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java b/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java
index 8ba5b99..c0089e5 100644
--- a/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java
+++ b/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java
@@ -226,6 +226,8 @@
                             if (!activeControl.getSurfacePosition().equals(lastSurfacePosition)
                                     && mAnimation != null) {
                                 startAnimation(mImeShowing, true /* forceRestart */);
+                            } else if (!mImeShowing) {
+                                removeImeSurface();
                             }
                         });
                     }
@@ -370,16 +372,7 @@
                     dispatchEndPositioning(mDisplayId, mCancelled, t);
                     if (mAnimationDirection == DIRECTION_HIDE && !mCancelled) {
                         t.hide(mImeSourceControl.getLeash());
-                        final IInputMethodManager imms = getImms();
-                        if (imms != null) {
-                            try {
-                                // Remove the IME surface to make the insets invisible for
-                                // non-client controlled insets.
-                                imms.removeImeSurface();
-                            } catch (RemoteException e) {
-                                Slog.e(TAG, "Failed to remove IME surface.", e);
-                            }
-                        }
+                        removeImeSurface();
                     }
                     t.apply();
                     mTransactionPool.release(t);
@@ -402,6 +395,19 @@
         }
     }
 
+    void removeImeSurface() {
+        final IInputMethodManager imms = getImms();
+        if (imms != null) {
+            try {
+                // Remove the IME surface to make the insets invisible for
+                // non-client controlled insets.
+                imms.removeImeSurface();
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to remove IME surface.", e);
+            }
+        }
+    }
+
     /**
      * Allows other things to synchronize with the ime position
      */
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
index 70f0399..05cf40a 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
@@ -120,6 +120,11 @@
     public abstract void reportImeControl(@Nullable IBinder windowToken);
 
     /**
+     * Destroys the IME surface.
+     */
+    public abstract void removeImeSurface();
+
+    /**
      * Fake implementation of {@link InputMethodManagerInternal}.  All the methods do nothing.
      */
     private static final InputMethodManagerInternal NOP =
@@ -166,6 +171,10 @@
                 @Override
                 public void reportImeControl(@Nullable IBinder windowToken) {
                 }
+
+                @Override
+                public void removeImeSurface() {
+                }
             };
 
     /**
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index c027ebc..d8ee32e 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -211,6 +211,7 @@
     static final int MSG_INITIALIZE_IME = 1040;
     static final int MSG_CREATE_SESSION = 1050;
     static final int MSG_REMOVE_IME_SURFACE = 1060;
+    static final int MSG_REMOVE_IME_SURFACE_FROM_WINDOW = 1061;
 
     static final int MSG_START_INPUT = 2000;
 
@@ -4005,6 +4006,13 @@
         mHandler.sendMessage(mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE));
     }
 
+    @Override
+    public void removeImeSurfaceFromWindow(IBinder windowToken) {
+        // No permission check, because we'll only execute the request if the calling window is
+        // also the current IME client.
+        mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE_FROM_WINDOW, windowToken).sendToTarget();
+    }
+
     @BinderThread
     private void notifyUserAction(@NonNull IBinder token) {
         if (DEBUG) {
@@ -4278,11 +4286,27 @@
                 return true;
             }
             case MSG_REMOVE_IME_SURFACE: {
-                try {
-                    if (mEnabledSession != null && mEnabledSession.session != null) {
-                        mEnabledSession.session.removeImeSurface();
+                synchronized (mMethodMap) {
+                    try {
+                        if (mEnabledSession != null && mEnabledSession.session != null
+                                && !mShowRequested) {
+                            mEnabledSession.session.removeImeSurface();
+                        }
+                    } catch (RemoteException e) {
                     }
-                } catch (RemoteException e) {
+                }
+                return true;
+            }
+            case MSG_REMOVE_IME_SURFACE_FROM_WINDOW: {
+                IBinder windowToken = (IBinder) msg.obj;
+                synchronized (mMethodMap) {
+                    try {
+                        if (windowToken == mCurFocusedWindow
+                                && mEnabledSession != null && mEnabledSession.session != null) {
+                            mEnabledSession.session.removeImeSurface();
+                        }
+                    } catch (RemoteException e) {
+                    }
                 }
                 return true;
             }
@@ -5116,6 +5140,11 @@
         public void reportImeControl(@Nullable IBinder windowToken) {
             mService.reportImeControl(windowToken);
         }
+
+        @Override
+        public void removeImeSurface() {
+            mService.mHandler.sendMessage(mService.mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE));
+        }
     }
 
     @BinderThread
diff --git a/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java
index 2e3d396..19dff98 100644
--- a/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java
@@ -225,6 +225,11 @@
                         @Override
                         public void reportImeControl(@Nullable IBinder windowToken) {
                         }
+
+                        @Override
+                        public void removeImeSurface() {
+                            reportNotSupported();
+                        }
                     });
         }
 
@@ -1473,6 +1478,12 @@
 
         @BinderThread
         @Override
+        public void removeImeSurfaceFromWindow(IBinder windowToken) {
+            reportNotSupported();
+        }
+
+        @BinderThread
+        @Override
         public boolean showSoftInput(
                 IInputMethodClient client, IBinder token, int flags,
                 ResultReceiver resultReceiver) {
diff --git a/services/core/java/com/android/server/pm/PackageInstallerService.java b/services/core/java/com/android/server/pm/PackageInstallerService.java
index 91b2ea1..08cdd8f 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerService.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerService.java
@@ -589,12 +589,12 @@
             }
         }
 
-        if (mBypassNextStagedInstallerCheck) {
-            mBypassNextStagedInstallerCheck = false;
-        } else if (params.isStaged
-                && !isCalledBySystemOrShell(callingUid)
-                && !isWhitelistedStagedInstaller(requestedInstallerPackageName)) {
-            throw new SecurityException("Installer not allowed to commit staged install");
+        if (params.isStaged && !isCalledBySystemOrShell(callingUid)) {
+            if (mBypassNextStagedInstallerCheck) {
+                mBypassNextStagedInstallerCheck = false;
+            } else if (!isStagedInstallerAllowed(requestedInstallerPackageName)) {
+                throw new SecurityException("Installer not allowed to commit staged install");
+            }
         }
 
         if (!params.isMultiPackage) {
@@ -725,7 +725,7 @@
                 || callingUid == Process.SHELL_UID;
     }
 
-    private boolean isWhitelistedStagedInstaller(String installerName) {
+    private boolean isStagedInstallerAllowed(String installerName) {
         return SystemConfig.getInstance().getWhitelistedStagedInstallers().contains(installerName);
     }
 
diff --git a/services/core/java/com/android/server/uri/UriGrantsManagerService.java b/services/core/java/com/android/server/uri/UriGrantsManagerService.java
index c38d649..4b3ddd8 100644
--- a/services/core/java/com/android/server/uri/UriGrantsManagerService.java
+++ b/services/core/java/com/android/server/uri/UriGrantsManagerService.java
@@ -114,7 +114,7 @@
     private static final boolean DEBUG = false;
     private static final String TAG = "UriGrantsManagerService";
     // Maximum number of persisted Uri grants a package is allowed
-    private static final int MAX_PERSISTED_URI_GRANTS = 128;
+    private static final int MAX_PERSISTED_URI_GRANTS = 512;
     private static final boolean ENABLE_DYNAMIC_PERMISSIONS = false;
 
     private final Object mLock = new Object();
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index e8a4234..b94fb04 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -3574,12 +3574,11 @@
         }
     }
 
-    private void updateImeControlTarget() {
+    void updateImeControlTarget() {
         mInputMethodControlTarget = computeImeControlTarget();
         mInsetsStateController.onImeControlTargetChanged(mInputMethodControlTarget);
 
-        final WindowState win = mInputMethodControlTarget != null
-                ? mInputMethodControlTarget.getWindow() : null;
+        final WindowState win = InsetsControlTarget.asWindowOrNull(mInputMethodControlTarget);
         final IBinder token = win != null ? win.mClient.asBinder() : null;
         // Note: not allowed to call into IMMS with the WM lock held, hence the post.
         mWmService.mH.post(() ->
@@ -3603,6 +3602,17 @@
         if (!isImeControlledByApp() && mRemoteInsetsControlTarget != null) {
             return mRemoteInsetsControlTarget;
         } else {
+            // Now, a special case -- if the last target's window is in the process of exiting, but
+            // not removed, keep on the last target to avoid IME flicker.
+            final WindowState cur = InsetsControlTarget.asWindowOrNull(mInputMethodControlTarget);
+            if (cur != null && !cur.mRemoved && cur.isDisplayedLw() && cur.isClosing()
+                    && !cur.isActivityTypeHome()) {
+                if (DEBUG_INPUT_METHOD) {
+                    Slog.v(TAG_WM, "Not changing control while current window"
+                            + " is closing and not removed");
+                }
+                return cur;
+            }
             // Otherwise, we just use the ime target as received from IME.
             return mInputMethodInputTarget;
         }
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index 8298763..99ee5e1 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -18,13 +18,14 @@
 
 import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_IME;
 
-import android.graphics.PixelFormat;
 import android.view.InsetsSource;
 import android.view.WindowInsets;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.protolog.common.ProtoLog;
 
+import java.io.PrintWriter;
+
 /**
  * Controller for IME inset source on the server. It's called provider as it provides the
  * {@link InsetsSource} to the client that uses it in {@link InsetsSourceConsumer}.
@@ -132,8 +133,17 @@
                 || (mImeTargetFromIme != null && dcTarget.getParentWindow() == mImeTargetFromIme
                         && dcTarget.mSubLayer > mImeTargetFromIme.mSubLayer)
                 || mImeTargetFromIme == mDisplayContent.getImeFallback()
-                // If IME target is transparent but control target matches requesting window.
-                || (controlTarget == mImeTargetFromIme
-                        && PixelFormat.formatHasAlpha(dcTarget.mAttrs.format));
+                || (!mImeTargetFromIme.isClosing() && controlTarget == mImeTargetFromIme);
+    }
+
+    @Override
+    public void dump(PrintWriter pw, String prefix) {
+        super.dump(pw, prefix);
+        if (mImeTargetFromIme != null) {
+            pw.print(prefix);
+            pw.print("showImePostLayout pending for mImeTargetFromIme=");
+            pw.print(mImeTargetFromIme);
+            pw.println();
+        }
     }
 }
diff --git a/services/core/java/com/android/server/wm/InsetsControlTarget.java b/services/core/java/com/android/server/wm/InsetsControlTarget.java
index c50f296..3ffc26a 100644
--- a/services/core/java/com/android/server/wm/InsetsControlTarget.java
+++ b/services/core/java/com/android/server/wm/InsetsControlTarget.java
@@ -62,4 +62,8 @@
         return false;
     }
 
+    /** Returns {@code target.getWindow()}, or null if {@code target} is {@code null}. */
+    static WindowState asWindowOrNull(InsetsControlTarget target) {
+        return target != null ? target.getWindow() : null;
+    }
 }
diff --git a/services/core/java/com/android/server/wm/InsetsStateController.java b/services/core/java/com/android/server/wm/InsetsStateController.java
index 63083fa..77bd4a4 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -44,6 +44,7 @@
 import android.view.InsetsState.InternalInsetsType;
 import android.view.WindowManager;
 
+import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.protolog.common.ProtoLog;
 
 import java.io.PrintWriter;
@@ -74,7 +75,21 @@
             w.notifyInsetsChanged();
         }
     };
-    private final InsetsControlTarget mEmptyImeControlTarget = new InsetsControlTarget() { };
+    private final InsetsControlTarget mEmptyImeControlTarget = new InsetsControlTarget() {
+        @Override
+        public void notifyInsetsControlChanged() {
+            InsetsSourceControl[] controls = getControlsForDispatch(this);
+            if (controls == null) {
+                return;
+            }
+            for (InsetsSourceControl control : controls) {
+                if (control.getType() == ITYPE_IME) {
+                    mDisplayContent.mWmService.mH.post(() ->
+                            InputMethodManagerInternal.get().removeImeSurface());
+                }
+            }
+        }
+    };
 
     InsetsStateController(DisplayContent displayContent) {
         mDisplayContent = displayContent;
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 837fafe..00be75f 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2170,6 +2170,9 @@
         if (isInputMethodTarget()) {
             dc.computeImeTarget(true /* updateImeTarget */);
         }
+        if (dc.mInputMethodControlTarget == this) {
+            dc.updateImeControlTarget();
+        }
 
         final int type = mAttrs.type;
         if (WindowManagerService.excludeWindowTypeFromTapOutTask(type)) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 77e3c59..d64fdb8 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -898,6 +898,26 @@
     }
 
     @Test
+    public void testComputeImeControlTarget_exitingApp() throws Exception {
+        final DisplayContent dc = createNewDisplay();
+
+        WindowState exitingWin = createWindow(null, TYPE_BASE_APPLICATION, "exiting app");
+        makeWindowVisible(exitingWin);
+        exitingWin.mWinAnimator.mDrawState = WindowStateAnimator.HAS_DRAWN;
+        exitingWin.mAnimatingExit = true;
+
+        dc.mInputMethodControlTarget = exitingWin;
+        dc.mInputMethodTarget = dc.mInputMethodInputTarget =
+                createWindow(null, TYPE_BASE_APPLICATION, "starting app");
+
+        assertEquals(exitingWin, dc.computeImeControlTarget());
+
+        exitingWin.removeImmediately();
+
+        assertEquals(dc.mInputMethodInputTarget, dc.computeImeControlTarget());
+    }
+
+    @Test
     public void testComputeImeControlTarget_splitscreen() throws Exception {
         final DisplayContent dc = createNewDisplay();
         dc.mInputMethodInputTarget = createWindow(null, TYPE_BASE_APPLICATION, "app");