diff --git a/core/java/android/view/ImeFocusController.java b/core/java/android/view/ImeFocusController.java
index f4daec121fc60..d4875d400c5b7 100644
--- a/core/java/android/view/ImeFocusController.java
+++ b/core/java/android/view/ImeFocusController.java
@@ -137,12 +137,6 @@ public final class ImeFocusController {
         boolean forceFocus = false;
         final InputMethodManagerDelegate immDelegate = getImmDelegate();
         synchronized (immDelegate.getLockObject()) {
-            // TODO(b/244504062): Remove isRestartOnNextWindowFocus.
-            if (immDelegate.isRestartOnNextWindowFocus(true /* reset */)) {
-                if (DEBUG) Log.v(TAG, "Restarting due to isRestartOnNextWindowFocus as true");
-                forceFocus = true;
-            }
-
             // Update mNextServedView when focusedView changed.
             onViewFocusChanged(viewForWindowFocus, true);
 
@@ -155,7 +149,7 @@ public final class ImeFocusController {
             }
         }
 
-        immDelegate.startInputAsyncOnWindowFocusGain(viewForWindowFocus,
+        immDelegate.startInputOnWindowFocusGain(viewForWindowFocus,
                 windowAttribute.softInputMode, windowAttribute.flags, forceFocus);
     }
 
@@ -321,7 +315,7 @@ public final class ImeFocusController {
          * {@link InputMethodManagerDelegate#getLockObject()} while {@link InputMethodManager}
          * calling into app-code in different threads.
          */
-        void startInputAsyncOnWindowFocusGain(View rootView,
+        void startInputOnWindowFocusGain(View rootView,
                 @WindowManager.LayoutParams.SoftInputModeFlags int softInputMode, int windowFlags,
                 boolean forceNewFocus);
         void finishInput();
@@ -330,7 +324,6 @@ public final class ImeFocusController {
         void finishComposingText();
         void setCurrentRootView(ViewRootImpl rootView);
         boolean isCurrentRootView(ViewRootImpl rootView);
-        boolean isRestartOnNextWindowFocus(boolean reset);
         boolean hasActiveConnection(View view);
 
         /**
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index b18ca86ff3667..7764526ae30df 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -771,10 +771,10 @@ public final class InputMethodManager {
         }
 
         /**
-         * For {@link ImeFocusController} to start input asynchronously when focus gain.
+         * For {@link ImeFocusController} to start input when gaining the window focus.
          */
         @Override
-        public void startInputAsyncOnWindowFocusGain(View focusedView,
+        public void startInputOnWindowFocusGain(View focusedView,
                 @SoftInputModeFlags int softInputMode, int windowFlags, boolean forceNewFocus) {
             int startInputFlags = getStartInputFlags(focusedView, 0);
             startInputFlags |= StartInputFlags.WINDOW_GAINED_FOCUS;
@@ -787,6 +787,15 @@ public final class InputMethodManager {
             if (controller == null) {
                 return;
             }
+
+            synchronized (mH) {
+                if (mRestartOnNextWindowFocus) {
+                    if (DEBUG) Log.v(TAG, "Restarting due to mRestartOnNextWindowFocus as true");
+                    mRestartOnNextWindowFocus = false;
+                    forceNewFocus = true;
+                }
+            }
+
             if (controller.checkFocus(forceNewFocus, false)) {
                 // We need to restart input on the current focus view.  This
                 // should be done in conjunction with telling the system service
@@ -859,20 +868,6 @@ public final class InputMethodManager {
             }
         }
 
-        /**
-         * For {@link ImeFocusController#checkFocus} if needed to force check new focus.
-         */
-        @Override
-        public boolean isRestartOnNextWindowFocus(boolean reset) {
-            synchronized (mH) {
-                final boolean result = mRestartOnNextWindowFocus;
-                if (reset) {
-                    mRestartOnNextWindowFocus = false;
-                }
-                return result;
-            }
-        }
-
         /**
          * Checks whether the active input connection (if any) is for the given view.
          *