diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index b1f027a..bef6258 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.taskbar;
 
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
+import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
 import static android.view.WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR_PANEL;
 
@@ -197,6 +198,7 @@
         mWindowLayoutParams.packageName = getPackageName();
         mWindowLayoutParams.gravity = Gravity.BOTTOM;
         mWindowLayoutParams.setFitInsetsTypes(0);
+        mWindowLayoutParams.receiveInsetsIgnoringZOrder = true;
         mWindowLayoutParams.softInputMode = WindowManager.LayoutParams.SOFT_INPUT_ADJUST_NOTHING;
         mWindowLayoutParams.layoutInDisplayCutoutMode = LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
         mWindowLayoutParams.privateFlags =
@@ -268,14 +270,6 @@
     }
 
     @Override
-    public boolean supportsIme() {
-        // Currently we don't support IME because we have FLAG_NOT_FOCUSABLE. We can remove that
-        // flag when opening a floating view that needs IME (such as Folder), but then that means
-        // Taskbar will be below IME and thus users can't click the back button.
-        return false;
-    }
-
-    @Override
     public View.OnClickListener getItemOnClickListener() {
         return this::onTaskbarIconClicked;
     }
@@ -501,6 +495,19 @@
         return mTaskbarHeightForIme;
     }
 
+    /**
+     * Either adds or removes {@link WindowManager.LayoutParams#FLAG_NOT_FOCUSABLE} on the taskbar
+     * window.
+     */
+    public void setTaskbarWindowFocusableForIme(boolean focusable) {
+        if (focusable) {
+            mWindowLayoutParams.flags &= ~FLAG_NOT_FOCUSABLE;
+        } else {
+            mWindowLayoutParams.flags |= FLAG_NOT_FOCUSABLE;
+        }
+        mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+    }
+
     protected void onTaskbarIconClicked(View view) {
         Object tag = view.getTag();
         if (tag instanceof Task) {
@@ -510,6 +517,17 @@
         } else if (tag instanceof FolderInfo) {
             FolderIcon folderIcon = (FolderIcon) view;
             Folder folder = folderIcon.getFolder();
+
+            folder.setOnFolderStateChangedListener(newState -> {
+                if (newState == Folder.STATE_OPEN) {
+                    setTaskbarWindowFocusableForIme(true);
+                } else if (newState == Folder.STATE_CLOSED) {
+                    // Defer by a frame to ensure we're no longer fullscreen and thus won't jump.
+                    getDragLayer().post(() -> setTaskbarWindowFocusableForIme(false));
+                    folder.setOnFolderStateChangedListener(null);
+                }
+            });
+
             setTaskbarWindowFullscreen(true);
 
             getDragLayer().post(() -> {
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index daef682..1e6342c 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -57,6 +57,7 @@
 import android.view.inputmethod.EditorInfo;
 import android.widget.TextView;
 
+import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
 import androidx.core.content.res.ResourcesCompat;
 
@@ -101,6 +102,8 @@
 import com.android.launcher3.views.ClipPathView;
 import com.android.launcher3.widget.PendingAddShortcutInfo;
 
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -130,10 +133,13 @@
      */
     private static final int MIN_CONTENT_DIMEN = 5;
 
-    static final int STATE_NONE = -1;
-    static final int STATE_SMALL = 0;
-    static final int STATE_ANIMATING = 1;
-    static final int STATE_OPEN = 2;
+    public static final int STATE_CLOSED = 0;
+    public static final int STATE_ANIMATING = 1;
+    public static final int STATE_OPEN = 2;
+
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef({STATE_CLOSED, STATE_ANIMATING, STATE_OPEN})
+    public @interface FolderState {}
 
     /**
      * Time for which the scroll hint is shown before automatically changing page.
@@ -198,13 +204,12 @@
 
     @ViewDebug.ExportedProperty(category = "launcher",
             mapping = {
-                    @ViewDebug.IntToString(from = STATE_NONE, to = "STATE_NONE"),
-                    @ViewDebug.IntToString(from = STATE_SMALL, to = "STATE_SMALL"),
+                    @ViewDebug.IntToString(from = STATE_CLOSED, to = "STATE_CLOSED"),
                     @ViewDebug.IntToString(from = STATE_ANIMATING, to = "STATE_ANIMATING"),
                     @ViewDebug.IntToString(from = STATE_OPEN, to = "STATE_OPEN"),
             })
-    @Thunk
-    int mState = STATE_NONE;
+    private int mState = STATE_CLOSED;
+    private OnFolderStateChangedListener mOnFolderStateChangedListener;
     @ViewDebug.ExportedProperty(category = "launcher")
     private boolean mRearrangeOnClose = false;
     boolean mItemsInvalidated = false;
@@ -277,19 +282,15 @@
         mPageIndicator = findViewById(R.id.folder_page_indicator);
         mFolderName = findViewById(R.id.folder_name);
         mFolderName.setTextSize(TypedValue.COMPLEX_UNIT_PX, dp.folderLabelTextSizePx);
-        if (mActivityContext.supportsIme()) {
-            mFolderName.setOnBackKeyListener(this);
-            mFolderName.setOnFocusChangeListener(this);
-            mFolderName.setOnEditorActionListener(this);
-            mFolderName.setSelectAllOnFocus(true);
-            mFolderName.setInputType(mFolderName.getInputType()
-                    & ~InputType.TYPE_TEXT_FLAG_AUTO_CORRECT
-                    | InputType.TYPE_TEXT_FLAG_NO_SUGGESTIONS
-                    | InputType.TYPE_TEXT_FLAG_CAP_WORDS);
-            mFolderName.forceDisableSuggestions(true);
-        } else {
-            mFolderName.setEnabled(false);
-        }
+        mFolderName.setOnBackKeyListener(this);
+        mFolderName.setOnFocusChangeListener(this);
+        mFolderName.setOnEditorActionListener(this);
+        mFolderName.setSelectAllOnFocus(true);
+        mFolderName.setInputType(mFolderName.getInputType()
+                & ~InputType.TYPE_TEXT_FLAG_AUTO_CORRECT
+                | InputType.TYPE_TEXT_FLAG_NO_SUGGESTIONS
+                | InputType.TYPE_TEXT_FLAG_CAP_WORDS);
+        mFolderName.forceDisableSuggestions(true);
 
         mFooter = findViewById(R.id.folder_footer);
         mFooterHeight = getResources().getDimensionPixelSize(R.dimen.folder_label_height);
@@ -561,7 +562,7 @@
         a.addListener(new AnimatorListenerAdapter() {
             @Override
             public void onAnimationStart(Animator animation) {
-                mState = STATE_ANIMATING;
+                setState(STATE_ANIMATING);
                 mCurrentAnimator = a;
             }
 
@@ -686,7 +687,7 @@
 
             @Override
             public void onAnimationEnd(Animator animation) {
-                mState = STATE_OPEN;
+                setState(STATE_OPEN);
                 announceAccessibilityChanges();
                 AccessibilityManagerCompat.sendFolderOpenedEventToTest(getContext());
 
@@ -862,7 +863,7 @@
         }
         mSuppressFolderDeletion = false;
         clearDragInfo();
-        mState = STATE_SMALL;
+        setState(STATE_CLOSED);
         mContent.setCurrentPage(0);
     }
 
@@ -1655,4 +1656,21 @@
 
         return windowBottomPx - folderBottomPx;
     }
+
+    private void setState(@FolderState int newState) {
+        mState = newState;
+        if (mOnFolderStateChangedListener != null) {
+            mOnFolderStateChangedListener.onFolderStateChanged(mState);
+        }
+    }
+
+    public void setOnFolderStateChangedListener(@Nullable OnFolderStateChangedListener listener) {
+        mOnFolderStateChangedListener = listener;
+    }
+
+    /** Listener that can be registered via {@link Folder#setOnFolderStateChangedListener} */
+    public interface OnFolderStateChangedListener {
+        /** See {@link Folder.FolderState} */
+        void onFolderStateChanged(@FolderState int newState);
+    }
 }
diff --git a/src/com/android/launcher3/util/UiThreadHelper.java b/src/com/android/launcher3/util/UiThreadHelper.java
index ac5368c..a1f31fe 100644
--- a/src/com/android/launcher3/util/UiThreadHelper.java
+++ b/src/com/android/launcher3/util/UiThreadHelper.java
@@ -28,7 +28,6 @@
 import android.view.View;
 import android.view.inputmethod.InputMethodManager;
 
-import com.android.launcher3.BaseActivity;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -56,7 +55,7 @@
                 STATS_LOGGER_KEY,
                 Message.obtain(
                         HANDLER.get(root.getContext()),
-                        () -> BaseActivity.fromContext(root.getContext())
+                        () -> ActivityContext.lookupContext(root.getContext())
                                 .getStatsLogManager()
                                 .logger()
                                 .log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED)
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index e09eff6..c0f6316 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -125,13 +125,6 @@
     }
 
     /**
-     * Returns whether we can show the IME for elements hosted by this ActivityContext.
-     */
-    default boolean supportsIme() {
-        return true;
-    }
-
-    /**
      * Called just before logging the given item.
      */
     default void applyOverwritesToLogItem(LauncherAtom.ItemInfo.Builder itemInfoBuilder) { }