diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index 02b2c5d5db84c..d609fb8eb2342 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -124,7 +124,12 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     public void setControl(@Nullable InsetsSourceControl control, int[] showTypes,
             int[] hideTypes) {
         super.setControl(control, showTypes, hideTypes);
-        if (control == null && !isRequestedVisibleAwaitingControl()) {
+        // TODO(b/204524304): clean-up how to deal with the timing issues of hiding IME:
+        //  1) Already requested show IME, in the meantime of WM callback the control but got null
+        //  control when relayout comes first
+        //  2) Make sure no regression on some implicit request IME visibility calls (e.g.
+        //  toggleSoftInput)
+        if (control == null && !mIsRequestedVisibleAwaitingControl) {
             hide();
             removeSurface();
         }