diff --git a/core/java/android/view/ImeFocusController.java b/core/java/android/view/ImeFocusController.java
index a5ac9480d2cfa..89e901cd70fe9 100644
--- a/core/java/android/view/ImeFocusController.java
+++ b/core/java/android/view/ImeFocusController.java
@@ -43,8 +43,21 @@ public final class ImeFocusController {
 
     private final ViewRootImpl mViewRootImpl;
     private boolean mHasImeFocus = false;
+
+    /**
+     * This is the view that should currently be served by an input method,
+     * regardless of the state of setting that up.
+     * @see InputMethodManagerDelegate#getLockObject()
+     */
     private View mServedView;
+
+    /**
+     * This is the next view that will be served by the input method, when
+     * we get around to updating things.
+     * @see InputMethodManagerDelegate#getLockObject()
+     */
     private View mNextServedView;
+
     private InputMethodManagerDelegate mDelegate;
 
     @UiThread
@@ -284,21 +297,44 @@ public final class ImeFocusController {
         boolean isCurrentRootView(ViewRootImpl rootView);
         boolean isRestartOnNextWindowFocus(boolean reset);
         boolean hasActiveConnection(View view);
+
+        /**
+         * Returns the {@code InputMethodManager#mH} lock object.
+         * Used for {@link ImeFocusController} to guard the served view being accessed by
+         * {@link InputMethodManager} in different threads.
+         *
+         * TODO(b/244504062): Use this to all places requires synchronization in controller.
+         */
+        Object getLockObject();
     }
 
-    public View getServedView() {
+    /**
+     * Returns The current IME served view for {@link InputMethodManager}.
+     * Used to start input connection or check the caller's validity when calling
+     * {@link InputMethodManager} APIs.
+     * Note that this method requires to be called inside {@code InputMethodManager#mH} lock for
+     * data consistency.
+     */
+    public View getServedViewLocked() {
         return mServedView;
     }
 
-    public View getNextServedView() {
+    /**
+     * Returns The next incoming IME served view for {@link InputMethodManager}.
+     * Note that this method requires to be called inside {@code InputMethodManager#mH} lock for
+     * data consistency.
+     */
+    public View getNextServedViewLocked() {
         return mNextServedView;
     }
 
-    public void setServedView(View view) {
+    // TODO(b/244504062): Remove this method dependency from InputMethodManager.
+    public void setServedViewLocked(View view) {
         mServedView = view;
     }
 
-    public void setNextServedView(View view) {
+    // TODO(b/244504062): Remove this method dependency from InputMethodManager.
+    public void setNextServedViewLocked(View view) {
         mNextServedView = view;
     }
 
diff --git a/core/java/android/view/inputmethod/BaseInputConnection.java b/core/java/android/view/inputmethod/BaseInputConnection.java
index 624937ff1d990..f0025922e1986 100644
--- a/core/java/android/view/inputmethod/BaseInputConnection.java
+++ b/core/java/android/view/inputmethod/BaseInputConnection.java
@@ -838,11 +838,8 @@ public class BaseInputConnection implements InputConnection {
             Context context;
             if (mTargetView != null) {
                 context = mTargetView.getContext();
-            } else if (mIMM.mCurRootView != null) {
-                final View servedView = mIMM.mCurRootView.getImeFocusController().getServedView();
-                context = servedView != null ? servedView.getContext() : null;
             } else {
-                context = null;
+                context = mIMM.getFallbackContextFromServedView();
             }
             if (context != null) {
                 TypedArray ta = context.getTheme()
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index 0c7c1639e6c73..84239385775ca 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -677,6 +677,21 @@ public final class InputMethodManager {
         return fallbackImm;
     }
 
+    /**
+     * An internal API that returns the {@link Context} of the current served view connected to
+     * an input method.
+     * @hide
+     */
+    Context getFallbackContextFromServedView() {
+        synchronized (mH) {
+            if (mCurRootView == null) {
+                return null;
+            }
+            final View servedView = mCurRootView.getImeFocusController().getServedViewLocked();
+            return servedView != null ? servedView.getContext() : null;
+        }
+    }
+
     private static boolean canStartInput(View servedView) {
         // We can start input ether the servedView has window focus
         // or the activity is showing autofill ui.
@@ -786,7 +801,7 @@ public final class InputMethodManager {
             synchronized (mH) {
                 // For some reason we didn't do a startInput + windowFocusGain, so
                 // we'll just do a window focus gain and call it a day.
-                View servedView = controller.getServedView();
+                View servedView = controller.getServedViewLocked();
                 boolean nextFocusHasConnection = servedView != null && servedView == focusedView
                         && hasActiveConnection(focusedView);
                 if (DEBUG) {
@@ -878,6 +893,16 @@ public final class InputMethodManager {
                         && mServedInputConnection.getServedView() == view;
             }
         }
+
+        /**
+         * Returns the {@link InputMethodManager#mH} lock object.
+         * Used for {@link ImeFocusController} to guard the served view being accessed by
+         * {@link InputMethodManager} in different threads.
+         */
+        @Override
+        public Object getLockObject() {
+            return mH;
+        }
     }
 
     /** @hide */
@@ -898,26 +923,27 @@ public final class InputMethodManager {
 
     @GuardedBy("mH")
     private View getServedViewLocked() {
-        return mCurRootView != null ? mCurRootView.getImeFocusController().getServedView() : null;
+        return mCurRootView != null ? mCurRootView.getImeFocusController().getServedViewLocked()
+                : null;
     }
 
     @GuardedBy("mH")
     private View getNextServedViewLocked() {
-        return mCurRootView != null ? mCurRootView.getImeFocusController().getNextServedView()
+        return mCurRootView != null ? mCurRootView.getImeFocusController().getNextServedViewLocked()
                 : null;
     }
 
     @GuardedBy("mH")
     private void setServedViewLocked(View view) {
         if (mCurRootView != null) {
-            mCurRootView.getImeFocusController().setServedView(view);
+            mCurRootView.getImeFocusController().setServedViewLocked(view);
         }
     }
 
     @GuardedBy("mH")
     private void setNextServedViewLocked(View view) {
         if (mCurRootView != null) {
-            mCurRootView.getImeFocusController().setNextServedView(view);
+            mCurRootView.getImeFocusController().setNextServedViewLocked(view);
         }
     }