diff --git a/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java b/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java
index e679f2655f64b..a2655f4c0f4d1 100644
--- a/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java
+++ b/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java
@@ -22,6 +22,7 @@ import static android.server.inputmethod.InputMethodManagerServiceProto.SHOW_EXP
 import static android.server.inputmethod.InputMethodManagerServiceProto.SHOW_FORCED;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.INVALID_DISPLAY;
+import static android.view.WindowManager.LayoutParams.SOFT_INPUT_IS_FORWARD_NAVIGATION;
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_UNSPECIFIED;
 import static android.view.WindowManager.LayoutParams.SoftInputModeFlags;
 
@@ -249,6 +250,29 @@ public final class ImeVisibilityStateComputer {
         return mService.mCurFocusedWindow;
     }
 
+    IBinder getWindowTokenFrom(ImeTargetWindowState windowState) {
+        for (IBinder windowToken : mRequestWindowStateMap.keySet()) {
+            final ImeTargetWindowState state = mRequestWindowStateMap.get(windowToken);
+            if (state == windowState) {
+                return windowToken;
+            }
+        }
+        return null;
+    }
+
+    boolean shouldRestoreImeVisibility(@NonNull ImeTargetWindowState state) {
+        final int softInputMode = state.getSoftInputModeState();
+        switch (softInputMode & WindowManager.LayoutParams.SOFT_INPUT_MASK_STATE) {
+            case WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_HIDDEN:
+                return false;
+            case WindowManager.LayoutParams.SOFT_INPUT_STATE_HIDDEN:
+                if ((softInputMode & SOFT_INPUT_IS_FORWARD_NAVIGATION) != 0) {
+                    return false;
+                }
+        }
+        return mWindowManagerInternal.shouldRestoreImeVisibility(getWindowTokenFrom(state));
+    }
+
     void dumpDebug(ProtoOutputStream proto, long fieldId) {
         proto.write(SHOW_EXPLICITLY_REQUESTED, mRequestedShowExplicitly);
         proto.write(SHOW_FORCED, mShowForced);
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 06ba880f66b68..2dbbb1085bb11 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -3666,8 +3666,9 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
 
         // Init the focused window state (e.g. whether the editor has focused or IME focus has
         // changed from another window).
-        mVisibilityStateComputer.setWindowState(windowToken,
-                new ImeTargetWindowState(softInputMode, !sameWindowFocused, isTextEditor));
+        final ImeTargetWindowState windowState = new ImeTargetWindowState(
+                softInputMode, !sameWindowFocused, isTextEditor);
+        mVisibilityStateComputer.setWindowState(windowToken, windowState);
 
         if (sameWindowFocused && isTextEditor) {
             if (DEBUG) {
@@ -3718,7 +3719,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         // Because the app might leverage these flags to hide soft-keyboard with showing their own
         // UI for input.
         if (isTextEditor && editorInfo != null
-                && shouldRestoreImeVisibility(windowToken, softInputMode)) {
+                && mVisibilityStateComputer.shouldRestoreImeVisibility(windowState)) {
             if (DEBUG) Slog.v(TAG, "Will show input to restore visibility");
             res = startInputUncheckedLocked(cs, inputContext, remoteAccessibilityInputConnection,
                     editorInfo, startInputFlags, startInputReason, unverifiedTargetSdkVersion,
@@ -3907,19 +3908,6 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         return true;
     }
 
-    private boolean shouldRestoreImeVisibility(IBinder windowToken,
-            @SoftInputModeFlags int softInputMode) {
-        switch (softInputMode & LayoutParams.SOFT_INPUT_MASK_STATE) {
-            case LayoutParams.SOFT_INPUT_STATE_ALWAYS_HIDDEN:
-                return false;
-            case LayoutParams.SOFT_INPUT_STATE_HIDDEN:
-                if ((softInputMode & LayoutParams.SOFT_INPUT_IS_FORWARD_NAVIGATION) != 0) {
-                    return false;
-                }
-        }
-        return mWindowManagerInternal.shouldRestoreImeVisibility(windowToken);
-    }
-
     @GuardedBy("ImfLock.class")
     private boolean canShowInputMethodPickerLocked(IInputMethodClient client) {
         final int uid = Binder.getCallingUid();