diff --git a/res/layout/widgets_search_bar.xml b/res/layout/widgets_search_bar.xml
index 252637d..450e5f1 100644
--- a/res/layout/widgets_search_bar.xml
+++ b/res/layout/widgets_search_bar.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="utf-8"?>
-<com.android.launcher3.widget.picker.search.WidgetsSearchBar
+<com.android.launcher3.widget.picker.search.LauncherWidgetsSearchBar
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:id="@+id/widgets_search_bar"
     android:layout_width="match_parent"
@@ -7,8 +7,7 @@
     android:orientation="horizontal"
     android:layout_marginTop="16dp"
     android:background="@drawable/bg_widgets_searchbox"
-    android:padding="12dp"
-    android:visibility="gone">
+    android:padding="12dp">
 
     <EditText
         android:id="@+id/widgets_search_bar_edit_text"
@@ -30,4 +29,4 @@
         android:background="?android:selectableItemBackground"
         android:layout_gravity="center"
         android:visibility="gone"/>
-</com.android.launcher3.widget.picker.search.WidgetsSearchBar>
\ No newline at end of file
+</com.android.launcher3.widget.picker.search.LauncherWidgetsSearchBar>
\ No newline at end of file
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 48e41d5..c1dd809 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -140,6 +140,9 @@
     public static final BooleanFlag ENABLE_OVERVIEW_SELECTIONS = new DeviceFlag(
             "ENABLE_OVERVIEW_SELECTIONS", true, "Show Select Mode button in Overview Actions");
 
+    public static final BooleanFlag ENABLE_WIDGETS_PICKER_AIAI_SEARCH = new DeviceFlag(
+            "ENABLE_WIDGETS_PICKER_AIAI_SEARCH", false, "Enable AiAi search in the widgets picker");
+
     public static final BooleanFlag ENABLE_OVERVIEW_SHARE = getDebugFlag(
             "ENABLE_OVERVIEW_SHARE", false, "Show Share button in Overview Actions");
 
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 6b3c71a..92994be 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -439,7 +439,8 @@
     public int getHeaderViewHeight() {
         return measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mCollapseHandle)
                 + measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mHeaderTitle)
-                + measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mSearchBar);
+                + measureHeightWithVerticalMargins(
+                (View) mSearchAndRecommendationViewHolder.mSearchBar);
     }
 
     /** private the height, in pixel, + the vertical margins of a given view. */
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
index 9ab6424..12d3f11 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
@@ -34,6 +34,7 @@
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
 /**
  * The widgets recycler view.
@@ -219,7 +220,8 @@
         int totalItemsHeight = 0;
         for (int i = 0; i < untilIndex; i++) {
             WidgetsListBaseEntry entry = mAdapter.getItems().get(i);
-            if (entry instanceof WidgetsListHeaderEntry) {
+            if (entry instanceof WidgetsListHeaderEntry ||
+                    entry instanceof WidgetsListSearchHeaderEntry) {
                 totalItemsHeight += mEstimatedWidgetListHeaderHeight;
             } else if (entry instanceof WidgetsListContentEntry) {
                 totalItemsHeight += mLastVisibleWidgetContentTableHeight;
diff --git a/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java b/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java
new file mode 100644
index 0000000..d68e87e
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget.picker.search;
+
+import android.content.Context;
+import android.util.AttributeSet;
+import android.widget.EditText;
+import android.widget.ImageButton;
+import android.widget.LinearLayout;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.R;
+import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+import java.util.List;
+
+/**
+ * View for a search bar with an edit text with a cancel button.
+ */
+public class LauncherWidgetsSearchBar extends LinearLayout implements WidgetsSearchBar {
+    private WidgetsSearchBarController mController;
+    private EditText mEditText;
+    private ImageButton mCancelButton;
+
+    public LauncherWidgetsSearchBar(Context context) {
+        this(context, null, 0);
+    }
+
+    public LauncherWidgetsSearchBar(@NonNull Context context,
+            @Nullable AttributeSet attrs) {
+        this(context, attrs, 0);
+    }
+
+    public LauncherWidgetsSearchBar(@NonNull Context context, @Nullable AttributeSet attrs,
+            int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+    }
+
+    @Override
+    public void initialize(List<WidgetsListBaseEntry> allWidgets,
+            SearchModeListener searchModeListener) {
+        SearchAlgorithm<WidgetsListBaseEntry> algo =
+                new SimpleWidgetsSearchAlgorithm(new SimpleWidgetsSearchPipeline(allWidgets));
+        mController = new WidgetsSearchBarController(
+                algo, mEditText, mCancelButton, searchModeListener);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mEditText = findViewById(R.id.widgets_search_bar_edit_text);
+        mCancelButton = findViewById(R.id.widgets_search_cancel_button);
+    }
+
+    @Override
+    protected void onDetachedFromWindow() {
+        super.onDetachedFromWindow();
+        mController.onDestroy();
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
index d8e9733..af6dc48 100644
--- a/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
+++ b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
@@ -16,64 +16,21 @@
 
 package com.android.launcher3.widget.picker.search;
 
-import android.content.Context;
-import android.util.AttributeSet;
-import android.widget.EditText;
-import android.widget.ImageButton;
-import android.widget.LinearLayout;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.R;
-import com.android.launcher3.search.SearchAlgorithm;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
 import java.util.List;
 
 /**
- * View for a search bar with an edit text with a cancel button.
+ * Interface for a widgets picker search bar.
  */
-public class WidgetsSearchBar extends LinearLayout {
-    private WidgetsSearchBarController mController;
-    private EditText mEditText;
-    private ImageButton mCancelButton;
-
-    public WidgetsSearchBar(Context context) {
-        this(context, null, 0);
-    }
-
-    public WidgetsSearchBar(@NonNull Context context,
-            @Nullable AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public WidgetsSearchBar(@NonNull Context context, @Nullable AttributeSet attrs,
-            int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-    }
-
+public interface WidgetsSearchBar {
     /**
      * Attaches a controller to the search bar which interacts with {@code searchModeListener}.
      */
-    public void initialize(List<WidgetsListBaseEntry> allWidgets,
-            SearchModeListener searchModeListener) {
-        SearchAlgorithm<WidgetsListBaseEntry> algo =
-                new SimpleWidgetsSearchAlgorithm(new SimpleWidgetsSearchPipeline(allWidgets));
-        mController = new WidgetsSearchBarController(
-                algo, mEditText, mCancelButton, searchModeListener);
-    }
+    void initialize(List<WidgetsListBaseEntry> allWidgets, SearchModeListener searchModeListener);
 
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        mEditText = findViewById(R.id.widgets_search_bar_edit_text);
-        mCancelButton = findViewById(R.id.widgets_search_cancel_button);
-    }
-
-    @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        mController.onDestroy();
-    }
+    /**
+     * Sets the vertical location, in pixels, of this search bar relative to its top position.
+     */
+    void setTranslationY(float translationY);
 }