Notify new IME target before IME starts interacting

Previously, there was a time window between when an IME starts new
input and when the IME issues an IPC IMM#setImeWindowStatus() so that
WindowManagerService (WMS) can be notified about the new IME target
window.

With this CL, it is now guaranteed that WindowManagerService (WMS) is
always notified about the new IME target window before IME starts
interacting with that window.

Note that WMS is not using notified IME target window yet hence there
should be no user-visible behavior change.

Bug: 110531072
Test: atest CtsInputMethodTestCases CtsInputMethodServiceHostTestCases
Change-Id: I032b91ce722a16b17518a5d88015c574d7d3e51b
diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index b4c8a5e..431c651 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -370,19 +370,6 @@
     InputConnection mStartedInputConnection;
     EditorInfo mInputEditorInfo;
 
-    /**
-     * A token to keep tracking the last IPC that triggered
-     * {@link #doStartInput(InputConnection, EditorInfo, boolean)}. If
-     * {@link #doStartInput(InputConnection, EditorInfo, boolean)} was not caused by IPCs from
-     * {@link com.android.server.InputMethodManagerService}, this needs to remain unchanged.
-     *
-     * <p>Some IPCs to {@link com.android.server.InputMethodManagerService} require this token to
-     * disentangle event flows for various purposes such as better window animation and providing
-     * fine-grained debugging information.</p>
-     */
-    @Nullable
-    private IBinder mStartInputToken;
-
     int mShowInputFlags;
     boolean mShowInputRequested;
     boolean mLastShowInputRequested;
@@ -528,7 +515,7 @@
         public void dispatchStartInputWithToken(@Nullable InputConnection inputConnection,
                 @NonNull EditorInfo editorInfo, boolean restarting,
                 @NonNull IBinder startInputToken) {
-            mStartInputToken = startInputToken;
+            mImm.reportStartInput(mToken, startInputToken);
 
             // This needs to be dispatched to interface methods rather than doStartInput().
             // Otherwise IME developers who have overridden those interface methods will lose
@@ -579,8 +566,8 @@
             }
             clearInsetOfPreviousIme();
             // If user uses hard keyboard, IME button should always be shown.
-            mImm.setImeWindowStatus(mToken, mStartInputToken,
-                    mapToImeWindowStatus(isInputViewShown()), mBackDisposition);
+            mImm.setImeWindowStatus(mToken, mapToImeWindowStatus(isInputViewShown()),
+                    mBackDisposition);
             if (resultReceiver != null) {
                 resultReceiver.send(wasVis != isInputViewShown()
                         ? InputMethodManager.RESULT_SHOWN
@@ -1059,8 +1046,8 @@
             }
             // If user uses hard keyboard, IME button should always be shown.
             boolean showing = onEvaluateInputViewShown();
-            mImm.setImeWindowStatus(mToken, mStartInputToken,
-                    IME_ACTIVE | (showing ? IME_VISIBLE : 0), mBackDisposition);
+            mImm.setImeWindowStatus(mToken, IME_ACTIVE | (showing ? IME_VISIBLE : 0),
+                    mBackDisposition);
         }
     }
 
@@ -1110,8 +1097,7 @@
             return;
         }
         mBackDisposition = disposition;
-        mImm.setImeWindowStatus(mToken, mStartInputToken, mapToImeWindowStatus(isInputViewShown()),
-                mBackDisposition);
+        mImm.setImeWindowStatus(mToken, mapToImeWindowStatus(isInputViewShown()), mBackDisposition);
     }
 
     /**
@@ -1861,8 +1847,7 @@
 
         final int nextImeWindowStatus = mapToImeWindowStatus(isInputViewShown());
         if (previousImeWindowStatus != nextImeWindowStatus) {
-            mImm.setImeWindowStatus(mToken, mStartInputToken, nextImeWindowStatus,
-                    mBackDisposition);
+            mImm.setImeWindowStatus(mToken, nextImeWindowStatus, mBackDisposition);
         }
         if ((previousImeWindowStatus & IME_ACTIVE) == 0) {
             if (DEBUG) Log.v(TAG, "showWindow: showing!");
@@ -1887,7 +1872,7 @@
     }
 
     private void doHideWindow() {
-        mImm.setImeWindowStatus(mToken, mStartInputToken, 0, mBackDisposition);
+        mImm.setImeWindowStatus(mToken, 0, mBackDisposition);
         hideWindow();
     }
 
@@ -2869,7 +2854,6 @@
         p.println("  mInputStarted=" + mInputStarted
                 + " mInputViewStarted=" + mInputViewStarted
                 + " mCandidatesViewStarted=" + mCandidatesViewStarted);
-        p.println("  mStartInputToken=" + mStartInputToken);
 
         if (mInputEditorInfo != null) {
             p.println("  mInputEditorInfo:");
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index 63ee624..bb93af5 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -798,10 +798,18 @@
     }
 
     /** @hide */
-    public void setImeWindowStatus(IBinder imeToken, IBinder startInputToken, int vis,
-            int backDisposition) {
+    public void setImeWindowStatus(IBinder imeToken, int vis, int backDisposition) {
         try {
-            mService.setImeWindowStatus(imeToken, startInputToken, vis, backDisposition);
+            mService.setImeWindowStatus(imeToken, vis, backDisposition);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /** @hide */
+    public void reportStartInput(IBinder imeToken, IBinder startInputToken) {
+        try {
+            mService.reportStartInput(imeToken, startInputToken);
         } 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 9ed1ffb..f76eddf 100644
--- a/core/java/com/android/internal/view/IInputMethodManager.aidl
+++ b/core/java/com/android/internal/view/IInputMethodManager.aidl
@@ -74,8 +74,8 @@
     void hideMySoftInput(in IBinder token, int flags);
     void showMySoftInput(in IBinder token, int flags);
     void updateStatusIcon(in IBinder token, String packageName, int iconId);
-    void setImeWindowStatus(in IBinder token, in IBinder startInputToken, int vis,
-            int backDisposition);
+    void setImeWindowStatus(in IBinder token, int vis, int backDisposition);
+    void reportStartInput(in IBinder token, in IBinder startInputToken);
     void registerSuggestionSpansForNotification(in SuggestionSpan[] spans);
     boolean notifySuggestionPicked(in SuggestionSpan span, String originalString, int index);
     InputMethodSubtype getCurrentInputMethodSubtype();
diff --git a/services/core/java/com/android/server/InputMethodManagerService.java b/services/core/java/com/android/server/InputMethodManagerService.java
index 20875a9..4b8ece9 100644
--- a/services/core/java/com/android/server/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/InputMethodManagerService.java
@@ -468,6 +468,13 @@
     IBinder mCurFocusedWindow;
 
     /**
+     * The last window token that we confirmed that IME started talking to.  This is always updated
+     * upon reports from the input method.  If the window state is already changed before the report
+     * is handled, this field just keeps the last value.
+     */
+    IBinder mLastImeTargetWindow;
+
+    /**
      * {@link WindowManager.LayoutParams#softInputMode} of {@link #mCurFocusedWindow}.
      *
      * @see #mCurFocusedWindow
@@ -688,7 +695,7 @@
     }
 
     @GuardedBy("mMethodMap")
-    private final WeakHashMap<IBinder, StartInputInfo> mStartInputMap = new WeakHashMap<>();
+    private final WeakHashMap<IBinder, IBinder> mImeTargetWindowMap = new WeakHashMap<>();
 
     /**
      * A ring buffer to store the history of {@link StartInputInfo}.
@@ -1809,7 +1816,7 @@
         final StartInputInfo info = new StartInputInfo(mCurToken, mCurId, startInputReason,
                 !initial, mCurFocusedWindow, mCurAttribute, mCurFocusedWindowSoftInputMode,
                 mCurSeq);
-        mStartInputMap.put(startInputToken, info);
+        mImeTargetWindowMap.put(startInputToken, mCurFocusedWindow);
         mStartInputHistory.addEntry(info);
 
         final SessionState session = mCurClient.curSession;
@@ -2219,15 +2226,12 @@
     @BinderThread
     @SuppressWarnings("deprecation")
     @Override
-    public void setImeWindowStatus(IBinder token, IBinder startInputToken, int vis,
-            int backDisposition) {
+    public void setImeWindowStatus(IBinder token, int vis, int backDisposition) {
         if (!calledWithValidToken(token)) {
             return;
         }
 
-        final StartInputInfo info;
         synchronized (mMethodMap) {
-            info = mStartInputMap.get(startInputToken);
             mImeWindowVis = vis;
             mBackDisposition = backDisposition;
             updateSystemUiLocked(token, vis, backDisposition);
@@ -2247,8 +2251,7 @@
                 break;
         }
         mWindowManagerInternal.updateInputMethodWindowStatus(token,
-                (vis & InputMethodService.IME_VISIBLE) != 0,
-                dismissImeOnBackKeyPressed, info != null ? info.mTargetWindow : null);
+                (vis & InputMethodService.IME_VISIBLE) != 0, dismissImeOnBackKeyPressed);
     }
 
     private void updateSystemUi(IBinder token, int vis, int backDisposition) {
@@ -2257,6 +2260,22 @@
         }
     }
 
+    @BinderThread
+    @Override
+    public void reportStartInput(IBinder token, IBinder startInputToken) {
+        if (!calledWithValidToken(token)) {
+            return;
+        }
+
+        synchronized (mMethodMap) {
+            final IBinder targetWindow = mImeTargetWindowMap.get(startInputToken);
+            if (targetWindow != null && mLastImeTargetWindow != targetWindow) {
+                mWindowManagerInternal.updateInputMethodTargetWindow(token, targetWindow);
+            }
+            mLastImeTargetWindow = targetWindow;
+        }
+    }
+
     // Caution! This method is called in this class. Handle multi-user carefully
     private void updateSystemUiLocked(IBinder token, int vis, int backDisposition) {
         if (!calledWithValidToken(token)) {
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index 6686b80..0c65518 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -368,14 +368,25 @@
      *                         to be synchronized with state in WindowManagerService.
      * @param dismissImeOnBackKeyPressed {@code true} if the software keyboard is shown and the back
      *                                   key is expected to dismiss the software keyboard.
-     * @param targetWindowToken token to identify the target window that the IME is associated with.
-     *                          {@code null} when application, system, or the IME itself decided to
-     *                          change its window visibility before being associated with any target
-     *                          window.
      */
     public abstract void updateInputMethodWindowStatus(@NonNull IBinder imeToken,
-            boolean imeWindowVisible, boolean dismissImeOnBackKeyPressed,
-            @Nullable IBinder targetWindowToken);
+            boolean imeWindowVisible, boolean dismissImeOnBackKeyPressed);
+
+    /**
+     * Notifies WindowManagerService that the current IME window status is being changed.
+     *
+     * <p>Only {@link com.android.server.InputMethodManagerService} is the expected and tested
+     * caller of this method.</p>
+     *
+     * @param imeToken token to track the active input method. Corresponding IME windows can be
+     *                 identified by checking {@link android.view.WindowManager.LayoutParams#token}.
+     *                 Note that there is no guarantee that the corresponding window is already
+     *                 created
+     * @param imeTargetWindowToken token to identify the target window that the IME is associated
+     *                             with
+     */
+    public abstract void updateInputMethodTargetWindow(@NonNull IBinder imeToken,
+            @NonNull IBinder imeTargetWindowToken);
 
     /**
       * Returns true when the hardware keyboard is available.
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 2c7ab7e..d43338d 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7364,16 +7364,18 @@
 
         @Override
         public void updateInputMethodWindowStatus(@NonNull IBinder imeToken,
-                boolean imeWindowVisible, boolean dismissImeOnBackKeyPressed,
-                @Nullable IBinder targetWindowToken) {
+                boolean imeWindowVisible, boolean dismissImeOnBackKeyPressed) {
+            mPolicy.setDismissImeOnBackKeyPressed(dismissImeOnBackKeyPressed);
+        }
+
+        @Override
+        public void updateInputMethodTargetWindow(@NonNull IBinder imeToken,
+                @NonNull IBinder imeTargetWindowToken) {
             // TODO (b/34628091): Use this method to address the window animation issue.
             if (DEBUG_INPUT_METHOD) {
-                Slog.w(TAG_WM, "updateInputMethodWindowStatus: imeToken=" + imeToken
-                        + " dismissImeOnBackKeyPressed=" + dismissImeOnBackKeyPressed
-                        + " imeWindowVisible=" + imeWindowVisible
-                        + " targetWindowToken=" + targetWindowToken);
+                Slog.w(TAG_WM, "updateInputMethodTargetWindow: imeToken=" + imeToken
+                        + " imeTargetWindowToken=" + imeTargetWindowToken);
             }
-            mPolicy.setDismissImeOnBackKeyPressed(dismissImeOnBackKeyPressed);
         }
 
         @Override