diff --git a/src/com/android/launcher3/ExtendedEditText.java b/src/com/android/launcher3/ExtendedEditText.java
index ec26f58..fe9348c 100644
--- a/src/com/android/launcher3/ExtendedEditText.java
+++ b/src/com/android/launcher3/ExtendedEditText.java
@@ -15,8 +15,6 @@
  */
 package com.android.launcher3;
 
-import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.SHOW;
-
 import android.content.Context;
 import android.graphics.Rect;
 import android.text.TextUtils;
@@ -93,7 +91,6 @@
      * @return true if the keyboard is shown correctly and focus is given to this view.
      */
     public boolean showKeyboard() {
-        onKeyboardShown();
         return requestFocus() && showSoftInputInternal();
     }
 
@@ -120,11 +117,6 @@
         }
     }
 
-    protected void onKeyboardShown() {
-        ActivityContext.lookupContext(getContext()).getStatsLogManager()
-                .keyboardStateManager().setKeyboardState(SHOW);
-    }
-
     private boolean showSoftInputInternal() {
         boolean result = false;
         InputMethodManager imm = getContext().getSystemService(InputMethodManager.class);
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 39b8de1..6b97c2c 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -19,6 +19,7 @@
 import static android.app.PendingIntent.FLAG_IMMUTABLE;
 import static android.app.PendingIntent.FLAG_UPDATE_CURRENT;
 import static android.content.pm.ActivityInfo.CONFIG_UI_MODE;
+import static android.view.WindowInsetsAnimation.Callback.DISPATCH_MODE_CONTINUE_ON_SUBTREE;
 import static android.view.accessibility.AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED;
 
 import static com.android.app.animation.Interpolators.EMPHASIZED;
@@ -67,6 +68,8 @@
 import static com.android.launcher3.config.FeatureFlags.ENABLE_SMARTSPACE_REMOVAL;
 import static com.android.launcher3.config.FeatureFlags.FOLDABLE_SINGLE_PAGE;
 import static com.android.launcher3.config.FeatureFlags.MULTI_SELECT_EDIT_MODE;
+import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.HIDE;
+import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.SHOW;
 import static com.android.launcher3.logging.StatsLogManager.EventEnum;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_HOME;
@@ -139,6 +142,8 @@
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.ViewTreeObserver.OnPreDrawListener;
+import android.view.WindowInsets;
+import android.view.WindowInsetsAnimation;
 import android.view.WindowManager.LayoutParams;
 import android.view.accessibility.AccessibilityEvent;
 import android.view.animation.OvershootInterpolator;
@@ -1084,6 +1089,25 @@
 
         DiscoveryBounce.showForHomeIfNeeded(this);
         mAppWidgetHolder.setActivityResumed(true);
+
+        // Listen for IME changes to keep state up to date.
+        getRootView().setWindowInsetsAnimationCallback(
+                new WindowInsetsAnimation.Callback(DISPATCH_MODE_CONTINUE_ON_SUBTREE) {
+                    @Override
+                    public WindowInsets onProgress(WindowInsets windowInsets,
+                            List<WindowInsetsAnimation> windowInsetsAnimations) {
+                        return windowInsets;
+                    }
+
+                    @Override
+                    public void onEnd(WindowInsetsAnimation animation) {
+                        WindowInsets insets = getRootView().getRootWindowInsets();
+                        boolean isImeVisible =
+                                insets != null && insets.isVisible(WindowInsets.Type.ime());
+                        getStatsLogManager().keyboardStateManager().setKeyboardState(
+                                isImeVisible ? SHOW : HIDE);
+                    }
+                });
     }
 
     private void logStopAndResume(boolean isResume) {
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index bef84f7..bf46bb0 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -19,7 +19,6 @@
 
 import static com.android.launcher3.LauncherSettings.Animation.DEFAULT_NO_ICON;
 import static com.android.launcher3.Utilities.allowBGLaunch;
-import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.HIDE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_KEYBOARD_CLOSED;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_PENDING_INTENT;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_TAP;
@@ -280,12 +279,9 @@
             if (wic != null) {
                 // Only hide the keyboard if it is actually showing.
                 if (isImeShown) {
-                    StatsLogManager slm = getStatsLogManager();
-                    slm.keyboardStateManager().setKeyboardState(HIDE);
-
                     // this method cannot be called cross threads
                     wic.hide(WindowInsets.Type.ime());
-                    slm.logger().log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED);
+                    getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED);
                 }
 
                 // If the WindowInsetsController is not null, we end here regardless of whether we