2 files changed, 12 insertions, 9 deletions
diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index 4fdea3b006dc..332e97c8bcf5 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -65,6 +65,9 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     public void onWindowFocusGained(boolean hasViewFocus) {
         super.onWindowFocusGained(hasViewFocus);
         getImm().registerImeConsumer(this);
+        if (isRequestedVisible() && getControl() == null) {
+            mIsRequestedVisibleAwaitingControl = true;
+        }
     }
 
     @Override
@@ -149,14 +152,11 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     }
 
     @Override
-    public void setControl(@Nullable InsetsSourceControl control, int[] showTypes,
+    public boolean setControl(@Nullable InsetsSourceControl control, int[] showTypes,
             int[] hideTypes) {
-        super.setControl(control, showTypes, hideTypes);
-        // 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 (!super.setControl(control, showTypes, hideTypes)) {
+            return false;
+        }
         if (control == null && !mIsRequestedVisibleAwaitingControl) {
             hide();
             removeSurface();
@@ -164,6 +164,7 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
         if (control != null) {
             mIsRequestedVisibleAwaitingControl = false;
         }
+        return true;
     }
 
     @Override
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index 6aab6359d23e..a82e19b10812 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -121,8 +121,9 @@ public class InsetsSourceConsumer {
      *                  animation should be run after setting the control.
      * @param hideTypes An integer array with a single entry that determines which types a hide
      *                  animation should be run after setting the control.
+     * @return Whether the control has changed from the server
      */
-    public void setControl(@Nullable InsetsSourceControl control,
+    public boolean setControl(@Nullable InsetsSourceControl control,
             @InsetsType int[] showTypes, @InsetsType int[] hideTypes) {
         if (mType == ITYPE_IME) {
             ImeTracing.getInstance().triggerClientDump("InsetsSourceConsumer#setControl",
@@ -133,7 +134,7 @@ public class InsetsSourceConsumer {
                 mSourceControl.release(SurfaceControl::release);
                 mSourceControl = control;
             }
-            return;
+            return false;
         }
         SurfaceControl oldLeash = mSourceControl != null ? mSourceControl.getLeash() : null;
 
@@ -198,6 +199,7 @@ public class InsetsSourceConsumer {
         if (lastControl != null) {
             lastControl.release(SurfaceControl::release);
         }
+        return true;
     }
 
     @VisibleForTesting