diff --git a/core/java/android/view/ImeFocusController.java b/core/java/android/view/ImeFocusController.java
index 89e901cd70fe9..2d2b0f5369d15 100644
--- a/core/java/android/view/ImeFocusController.java
+++ b/core/java/android/view/ImeFocusController.java
@@ -328,14 +328,23 @@ public final class ImeFocusController {
         return mNextServedView;
     }
 
-    // TODO(b/244504062): Remove this method dependency from InputMethodManager.
-    public void setServedViewLocked(View view) {
-        mServedView = view;
-    }
-
-    // TODO(b/244504062): Remove this method dependency from InputMethodManager.
-    public void setNextServedViewLocked(View view) {
-        mNextServedView = view;
+    /**
+     * Clears the served & the next served view when the controller triggers
+     * {@link InputMethodManagerDelegate#finishInput()} or
+     * {@link InputMethodManagerDelegate#finishInputAndReportToIme()}.
+     * Note that this method requires to be called inside {@code InputMethodManager#mH} lock for
+     * data consistency.
+     *
+     * @return The {@code mServedView} that has cleared, or {@code null} means nothing to clear.
+     */
+    public View clearServedViewsLocked() {
+        View clearedView = null;
+        mNextServedView = null;
+        if (mServedView != null) {
+            clearedView = mServedView;
+            mServedView = null;
+        }
+        return clearedView;
     }
 
     /**
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index 84239385775ca..b18ca86ff3667 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -933,20 +933,6 @@ public final class InputMethodManager {
                 : null;
     }
 
-    @GuardedBy("mH")
-    private void setServedViewLocked(View view) {
-        if (mCurRootView != null) {
-            mCurRootView.getImeFocusController().setServedViewLocked(view);
-        }
-    }
-
-    @GuardedBy("mH")
-    private void setNextServedViewLocked(View view) {
-        if (mCurRootView != null) {
-            mCurRootView.getImeFocusController().setNextServedViewLocked(view);
-        }
-    }
-
     private ImeFocusController getFocusController() {
         synchronized (mH) {
             if (mCurRootView != null) {
@@ -1793,13 +1779,13 @@ public final class InputMethodManager {
     @GuardedBy("mH")
     void finishInputLocked() {
         mVirtualDisplayToScreenMatrix = null;
-        setNextServedViewLocked(null);
-        if (getServedViewLocked() != null) {
+        final ImeFocusController controller = getFocusController();
+        final View clearedView = controller != null ? controller.clearServedViewsLocked() : null;
+        if (clearedView != null) {
             if (DEBUG) {
                 Log.v(TAG, "FINISH INPUT: mServedView="
-                        + InputMethodDebug.dumpViewInfo(getServedViewLocked()));
+                        + InputMethodDebug.dumpViewInfo(clearedView));
             }
-            setServedViewLocked(null);
             mCompletions = null;
             mServedConnecting = false;
             clearConnectionLocked();