diff --git a/Android.mk b/Android.mk
index 19ad328..127df79 100644
--- a/Android.mk
+++ b/Android.mk
@@ -145,7 +145,9 @@
     $(call all-java-files-under, quickstep/src) \
     $(call all-java-files-under, src_shortcuts_overrides)
 
-LOCAL_RESOURCE_DIR := $(LOCAL_PATH)/quickstep/res
+LOCAL_RESOURCE_DIR := \
+    $(LOCAL_PATH)/quickstep/res \
+    $(LOCAL_PATH)/quickstep/overview_ui_overrides/res
 LOCAL_PROGUARD_ENABLED := disabled
 
 
@@ -174,7 +176,9 @@
 LOCAL_OVERRIDES_PACKAGES := Home Launcher2 Launcher3
 LOCAL_REQUIRED_MODULES := privapp_whitelist_com.android.launcher3
 
-LOCAL_RESOURCE_DIR := $(LOCAL_PATH)/quickstep/res
+LOCAL_RESOURCE_DIR := \
+    $(LOCAL_PATH)/quickstep/res \
+    $(LOCAL_PATH)/quickstep/overview_ui_overrides/res
 
 LOCAL_FULL_LIBS_MANIFEST_FILES := \
     $(LOCAL_PATH)/quickstep/AndroidManifest-launcher.xml \
@@ -213,7 +217,8 @@
 LOCAL_RESOURCE_DIR := \
     $(LOCAL_PATH)/quickstep/res \
     $(LOCAL_PATH)/go/res \
-    $(LOCAL_PATH)/go/quickstep/res
+    $(LOCAL_PATH)/go/quickstep/res \
+    $(LOCAL_PATH)/go/quickstep/overview_ui_overrides/res
 
 LOCAL_PROGUARD_FLAG_FILES := proguard.flags
 LOCAL_PROGUARD_ENABLED := full
diff --git a/go/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml b/go/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
new file mode 100644
index 0000000..b438da3
--- /dev/null
+++ b/go/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
@@ -0,0 +1,98 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     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.
+-->
+<com.android.quickstep.views.GoOverviewActionsView
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="@dimen/overview_actions_height"
+    android:layout_gravity="center_horizontal|bottom">
+
+    <LinearLayout
+        android:id="@+id/action_buttons"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center"
+        android:orientation="horizontal">
+
+        <Space
+            android:layout_width="0dp"
+            android:layout_height="1dp"
+            android:layout_weight="1" />
+
+        <Button
+            android:id="@+id/action_listen"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_listen"
+            android:drawablePadding="1dp"
+            android:text="@string/action_listen"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Button
+            android:id="@+id/action_translate"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_translate"
+            android:drawablePadding="1dp"
+            android:text="@string/action_translate"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Button
+            android:id="@+id/action_search"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_search"
+            android:drawablePadding="1dp"
+            android:text="@string/action_search"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Button
+            android:id="@+id/action_screenshot"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_screenshot"
+            android:drawablePadding="1dp"
+            android:text="@string/action_screenshot"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Space
+            android:layout_width="0dp"
+            android:layout_height="1dp"
+            android:layout_weight="1" />
+
+        <Button
+            android:id="@+id/action_share"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableStart="@drawable/ic_share"
+            android:text="@string/action_share"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor"
+            android:visibility="gone" />
+
+        <Space
+            android:id="@+id/oav_three_button_space"
+            android:layout_width="0dp"
+            android:layout_height="1dp"
+            android:layout_weight="1"
+            android:visibility="gone" />
+    </LinearLayout>
+
+</com.android.quickstep.views.GoOverviewActionsView>
\ No newline at end of file
diff --git a/go/quickstep/overview_ui_overrides/res/values/config.xml b/go/quickstep/overview_ui_overrides/res/values/config.xml
new file mode 100644
index 0000000..ec21a01
--- /dev/null
+++ b/go/quickstep/overview_ui_overrides/res/values/config.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+/*
+* 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.
+*/
+-->
+<resources>
+    <string name="task_overlay_factory_class" translatable="false">
+        com.android.quickstep.TaskOverlayFactoryGo</string>
+</resources>
\ No newline at end of file
diff --git a/go/quickstep/res/drawable/ic_listen.xml b/go/quickstep/res/drawable/ic_listen.xml
new file mode 100644
index 0000000..a8e6c93
--- /dev/null
+++ b/go/quickstep/res/drawable/ic_listen.xml
@@ -0,0 +1,32 @@
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="28dp"
+    android:height="28dp"
+    android:viewportWidth="28"
+    android:viewportHeight="28">
+  <path
+      android:pathData="M10.5,15.17c2.58,0 4.67,-2.09 4.67,-4.67s-2.09,-4.67 -4.67,-4.67c-2.58,0 -4.67,2.09 -4.67,4.67S7.92,15.17 10.5,15.17zM10.5,8.17c1.28,0 2.33,1.05 2.33,2.33s-1.05,2.33 -2.33,2.33c-1.28,0 -2.33,-1.05 -2.33,-2.33S9.22,8.17 10.5,8.17z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M10.5,17.5c-3.11,0 -9.33,1.56 -9.33,4.67v2.33h18.67v-2.33C19.83,19.06 13.62,17.5 10.5,17.5zM3.5,22.17c0.26,-0.84 3.86,-2.33 7,-2.33c3.15,0 6.77,1.5 7,2.33H3.5z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M25.67,10.5c0,0.36 -0.02,0.71 -0.05,1.05c-0.01,0.15 -0.03,0.29 -0.05,0.43c-0.02,0.18 -0.05,0.36 -0.08,0.54c-0.04,0.2 -0.07,0.39 -0.12,0.58c-0.01,0.06 -0.03,0.11 -0.04,0.17c-0.59,2.34 -1.81,4.01 -2.52,4.82c-0.09,0.1 -0.18,0.2 -0.28,0.3c-0.17,0.18 -0.27,0.27 -0.27,0.27l-1.65,-1.63c1.34,-1.33 2.27,-3.07 2.6,-5.01c0.01,-0.08 0.02,-0.16 0.04,-0.24c0.06,-0.42 0.1,-0.85 0.1,-1.29c0,-0.44 -0.04,-0.88 -0.1,-1.3c-0.01,-0.06 -0.02,-0.13 -0.03,-0.19c-0.32,-1.95 -1.25,-3.7 -2.6,-5.04l1.65,-1.63c0,0 0.11,0.1 0.27,0.27c0.09,0.1 0.19,0.2 0.28,0.3c0.71,0.82 1.93,2.48 2.52,4.82c0.01,0.06 0.03,0.11 0.04,0.17c0.04,0.19 0.08,0.38 0.12,0.58c0.03,0.18 0.06,0.36 0.08,0.54c0.02,0.14 0.04,0.28 0.05,0.43C25.65,9.79 25.67,10.14 25.67,10.5z"
+      android:fillColor="#EA4335"/>
+  <path
+      android:pathData="M20.61,8.4C20.85,9.06 21,9.76 21,10.5s-0.15,1.44 -0.39,2.1c-0.28,0.77 -0.71,1.46 -1.25,2.05l-1.66,-1.64c0.56,-0.63 0.91,-1.44 0.95,-2.34c0,-0.06 0.02,-0.11 0.02,-0.17s-0.01,-0.11 -0.02,-0.17c-0.04,-0.9 -0.39,-1.71 -0.95,-2.34l1.66,-1.64C19.9,6.94 20.32,7.63 20.61,8.4z"
+      android:fillColor="#FBBC04"/>
+</vector>
diff --git a/go/quickstep/res/drawable/ic_search.xml b/go/quickstep/res/drawable/ic_search.xml
new file mode 100644
index 0000000..4307330
--- /dev/null
+++ b/go/quickstep/res/drawable/ic_search.xml
@@ -0,0 +1,32 @@
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="28dp"
+    android:height="28dp"
+    android:viewportWidth="28"
+    android:viewportHeight="28">
+  <path
+      android:pathData="M24.5,22.75l-6.84,-6.84c1,-1.35 1.59,-3.02 1.59,-4.83h-2.33c0,3.22 -2.62,5.83 -5.83,5.83v2.33c1.81,0 3.47,-0.6 4.83,-1.59l6.84,6.84L24.5,22.75z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M11.08,2.92v2.33c3.22,0 5.83,2.62 5.83,5.83h2.33C19.25,6.57 15.59,2.92 11.08,2.92z"
+      android:fillColor="#34A853"/>
+  <path
+      android:pathData="M5.25,11.08H2.92c0,4.51 3.66,8.17 8.17,8.17v-2.33C7.87,16.92 5.25,14.3 5.25,11.08z"
+      android:fillColor="#EA4335"/>
+  <path
+      android:pathData="M2.92,11.08h2.33c0,-3.22 2.62,-5.83 5.83,-5.83V2.92C6.57,2.92 2.92,6.57 2.92,11.08z"
+      android:fillColor="#FBBC04"/>
+</vector>
diff --git a/go/quickstep/res/drawable/ic_translate.xml b/go/quickstep/res/drawable/ic_translate.xml
new file mode 100644
index 0000000..1247807
--- /dev/null
+++ b/go/quickstep/res/drawable/ic_translate.xml
@@ -0,0 +1,32 @@
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="28dp"
+    android:height="28dp"
+    android:viewportWidth="28"
+    android:viewportHeight="28">
+  <path
+      android:pathData="M12.28,15.19l-0.07,-0.05c-0.61,-0.49 -1.15,-1.05 -1.65,-1.63c-1.05,-1.22 -1.88,-2.63 -2.39,-4.17H5.83c0.54,2.17 1.58,4.16 3.01,5.85l-5.93,5.23l1.75,1.75l5.91,-5.26c0.05,0.04 0.1,0.09 0.15,0.13l3.42,2.79l1.02,-2.33L12.28,15.19z"
+      android:fillColor="#FBBC04"/>
+  <path
+      android:pathData="M21.58,11.67h-2.33l-5.25,14h2.33l1.31,-3.5h5.54l1.32,3.5h2.33L21.58,11.67zM18.53,19.83l1.89,-5.05l1.89,5.05H18.53z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M11.67,2.33l-2.34,0l0,2.34l-8.16,0l0,2.33l10.5,0l0,-2.33z"
+      android:fillColor="#EA4335"/>
+  <path
+      android:pathData="M11.67,4.67V7H14c-0.61,2.42 -1.79,4.65 -3.44,6.5c0.5,0.59 1.04,1.15 1.65,1.63l0.07,0.05c2.03,-2.32 3.44,-5.14 4.05,-8.19h3.5V4.67H11.67z"
+      android:fillColor="#34A853"/>
+</vector>
diff --git a/go/quickstep/res/values/config.xml b/go/quickstep/res/values/config.xml
index f376774..9dca137 100644
--- a/go/quickstep/res/values/config.xml
+++ b/go/quickstep/res/values/config.xml
@@ -14,5 +14,11 @@
      limitations under the License.
 -->
 <resources>
+    <!-- The component to receive app sharing Intents -->
     <string name="app_sharing_component" translatable="false"/>
+    <!-- The package to receive Listen, Translate, and Search Intents -->
+    <string name="niu_actions_package" translatable="false"/>
+
+    <!-- Feature Flags -->
+    <bool name="enable_niu_actions">false</bool>
 </resources>
\ No newline at end of file
diff --git a/go/quickstep/res/values/strings.xml b/go/quickstep/res/values/strings.xml
index fdd8397..71e2f3a 100644
--- a/go/quickstep/res/values/strings.xml
+++ b/go/quickstep/res/values/strings.xml
@@ -3,4 +3,12 @@
 <resources xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
     <!-- Label for app share drop target. [CHAR_LIMIT=20] -->
     <string name="app_share_drop_target_label">Share App</string>
+
+    <!-- ******* Overview ******* -->
+    <!-- Label for a button that lets the user listen to the content of the current app. [CHAR_LIMIT=40] -->
+    <string name="action_listen">Listen</string>
+    <!-- Label for a button that translates a screenshot of the current app. [CHAR_LIMIT=40] -->
+    <string name="action_translate">Translate</string>
+    <!-- Label for a button that triggers Search on a screenshot of the current app. [CHAR_LIMIT=40] -->
+    <string name="action_search">Search</string>
 </resources>
diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
new file mode 100644
index 0000000..b102a39
--- /dev/null
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -0,0 +1,147 @@
+/*
+ * 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.quickstep;
+
+import static com.android.quickstep.views.OverviewActionsView.DISABLED_NO_THUMBNAIL;
+import static com.android.quickstep.views.OverviewActionsView.DISABLED_ROTATED;
+
+import android.annotation.SuppressLint;
+import android.content.Context;
+import android.content.Intent;
+import android.graphics.Matrix;
+import android.os.SystemClock;
+import android.text.TextUtils;
+
+import com.android.launcher3.R;
+import com.android.quickstep.views.OverviewActionsView;
+import com.android.quickstep.views.TaskThumbnailView;
+import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.ThumbnailData;
+
+/**
+ * Go-specific extension of the factory class that adds an overlay to TaskView
+ */
+public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
+    public static final String ACTION_LISTEN = "com.android.quickstep.ACTION_LISTEN";
+    public static final String ACTION_TRANSLATE = "com.android.quickstep.ACTION_TRANSLATE";
+    public static final String ACTION_SEARCH = "com.android.quickstep.ACTION_SEARCH";
+    public static final String ELAPSED_NANOS = "niu_actions_elapsed_realtime_nanos";
+
+    // Empty constructor required for ResourceBasedOverride
+    public TaskOverlayFactoryGo(Context context) {}
+
+    /**
+     * Create a new overlay instance for the given View
+     */
+    public TaskOverlayGo createOverlay(TaskThumbnailView thumbnailView) {
+        return new TaskOverlayGo(thumbnailView);
+    }
+
+    /**
+     * Overlay on each task handling Overview Action Buttons.
+     * @param <T> The type of View in which the overlay will be placed
+     */
+    public static final class TaskOverlayGo<T extends OverviewActionsView> extends TaskOverlay {
+
+        private String mPackageName;
+
+        private TaskOverlayGo(TaskThumbnailView taskThumbnailView) {
+            super(taskThumbnailView);
+        }
+
+        /**
+         * Called when the current task is interactive for the user
+         */
+        @Override
+        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
+                boolean rotated) {
+            getActionsView().updateDisabledFlags(DISABLED_NO_THUMBNAIL, thumbnail == null);
+            mPackageName =
+                    mApplicationContext.getResources().getString(R.string.niu_actions_package);
+
+            if (thumbnail == null || TextUtils.isEmpty(mPackageName)) {
+                return;
+            }
+
+            getActionsView().updateDisabledFlags(DISABLED_ROTATED, rotated);
+            boolean isAllowedByPolicy = thumbnail.isRealSnapshot;
+            getActionsView().setCallbacks(new OverlayUICallbacksGoImpl(isAllowedByPolicy, task));
+        }
+
+        private void sendNIUIntent(String actionType) {
+            Intent intent = createNIUIntent(actionType);
+            mImageApi.shareAsDataWithExplicitIntent(/* crop */ null, intent);
+        }
+
+        private Intent createNIUIntent(String actionType) {
+            return new Intent(actionType)
+                    .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK)
+                    .addFlags(Intent.FLAG_ACTIVITY_NO_ANIMATION)
+                    .setPackage(mPackageName)
+                    .putExtra(ELAPSED_NANOS, SystemClock.elapsedRealtimeNanos());
+        }
+
+        protected class OverlayUICallbacksGoImpl extends OverlayUICallbacksImpl
+                implements OverlayUICallbacksGo {
+            public OverlayUICallbacksGoImpl(boolean isAllowedByPolicy, Task task) {
+                super(isAllowedByPolicy, task);
+            }
+
+            @SuppressLint("NewApi")
+            public void onListen() {
+                if (mIsAllowedByPolicy) {
+                    sendNIUIntent(ACTION_LISTEN);
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+
+            @SuppressLint("NewApi")
+            public void onTranslate() {
+                if (mIsAllowedByPolicy) {
+                    sendNIUIntent(ACTION_TRANSLATE);
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+
+            @SuppressLint("NewApi")
+            public void onSearch() {
+                if (mIsAllowedByPolicy) {
+                    sendNIUIntent(ACTION_SEARCH);
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+        }
+    }
+
+    /**
+     * Callbacks the Ui can generate. This is the only way for a Ui to call methods on the
+     * controller.
+     */
+    public interface OverlayUICallbacksGo extends OverlayUICallbacks {
+        /** User has requested to listen to the current content read aloud */
+        void onListen();
+
+        /** User has requested a translation of the current content */
+        void onTranslate();
+
+        /** User has requested a visual search of the current content */
+        void onSearch();
+    }
+}
diff --git a/go/quickstep/src/com/android/quickstep/views/GoOverviewActionsView.java b/go/quickstep/src/com/android/quickstep/views/GoOverviewActionsView.java
new file mode 100644
index 0000000..9997d16
--- /dev/null
+++ b/go/quickstep/src/com/android/quickstep/views/GoOverviewActionsView.java
@@ -0,0 +1,75 @@
+/*
+ * 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.quickstep.views;
+
+import android.content.Context;
+import android.util.AttributeSet;
+import android.view.View;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.R;
+import com.android.quickstep.TaskOverlayFactoryGo.OverlayUICallbacksGo;
+
+/**
+ * View for showing Go-specific action buttons in Overview
+ */
+public final class GoOverviewActionsView extends OverviewActionsView<OverlayUICallbacksGo> {
+    public GoOverviewActionsView(Context context) {
+        this(context, null);
+    }
+
+    public GoOverviewActionsView(Context context, @Nullable AttributeSet attrs) {
+        this(context, attrs, 0);
+    }
+
+    public GoOverviewActionsView(Context context, @Nullable AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+
+        if (getResources().getBoolean(R.bool.enable_niu_actions)) {
+            findViewById(R.id.action_listen).setOnClickListener(this);
+            findViewById(R.id.action_translate).setOnClickListener(this);
+            findViewById(R.id.action_search).setOnClickListener(this);
+        } else {
+            findViewById(R.id.action_listen).setVisibility(View.GONE);
+            findViewById(R.id.action_translate).setVisibility(View.GONE);
+            findViewById(R.id.action_search).setVisibility(View.GONE);
+        }
+    }
+
+    @Override
+    public void onClick(View view) {
+        super.onClick(view);
+
+        if (mCallbacks == null) {
+            return;
+        }
+        int id = view.getId();
+        if (id == R.id.action_listen) {
+            mCallbacks.onListen();
+        } else if (id == R.id.action_translate) {
+            mCallbacks.onTranslate();
+        } else if (id == R.id.action_search) {
+            mCallbacks.onSearch();
+        }
+    }
+}
diff --git a/quickstep/res/layout/overview_actions_container.xml b/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
similarity index 100%
rename from quickstep/res/layout/overview_actions_container.xml
rename to quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
diff --git a/quickstep/overview_ui_overrides/res/values/config.xml b/quickstep/overview_ui_overrides/res/values/config.xml
new file mode 100644
index 0000000..0f09439
--- /dev/null
+++ b/quickstep/overview_ui_overrides/res/values/config.xml
@@ -0,0 +1,18 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<resources>
+    <string name="task_overlay_factory_class" translatable="false"/>
+</resources>
\ No newline at end of file
diff --git a/quickstep/res/layout/task.xml b/quickstep/res/layout/task.xml
index c90706b..0f9a6aa 100644
--- a/quickstep/res/layout/task.xml
+++ b/quickstep/res/layout/task.xml
@@ -19,7 +19,6 @@
     android:layout_height="match_parent"
     android:clipChildren="false"
     android:defaultFocusHighlightEnabled="false"
-    android:elevation="4dp"
     android:focusable="true">
 
     <com.android.quickstep.views.TaskThumbnailView
diff --git a/quickstep/res/layout/task_menu.xml b/quickstep/res/layout/task_menu.xml
index 744a305..3916ff9 100644
--- a/quickstep/res/layout/task_menu.xml
+++ b/quickstep/res/layout/task_menu.xml
@@ -29,6 +29,7 @@
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:gravity="center_horizontal"
+        android:layout_marginTop="16dp"
         android:layout_marginBottom="16dp"
         android:textSize="12sp"/>
 
diff --git a/quickstep/res/values/config.xml b/quickstep/res/values/config.xml
index 9ec303a..be66104 100644
--- a/quickstep/res/values/config.xml
+++ b/quickstep/res/values/config.xml
@@ -14,8 +14,6 @@
      limitations under the License.
 -->
 <resources>
-    <string name="task_overlay_factory_class" translatable="false"/>
-
     <string name="overscroll_plugin_factory_class" translatable="false" />
 
     <!-- Activities which block home gesture -->
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 749b4b2..68c3851 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -29,7 +29,7 @@
     <dimen name="overview_actions_bottom_margin_three_button">8dp</dimen>
     <dimen name="overview_actions_horizontal_margin">16dp</dimen>
 
-    <dimen name="recents_page_spacing">10dp</dimen>
+    <dimen name="recents_page_spacing">16dp</dimen>
     <dimen name="recents_clear_all_deadzone_vertical_margin">70dp</dimen>
 
     <!-- The speed in dp/s at which the user needs to be scrolling in recents such that we start
@@ -58,7 +58,7 @@
     <dimen name="task_card_menu_option_vertical_padding">8dp</dimen>
     <dimen name="task_card_menu_shadow_height">3dp</dimen>
     <dimen name="task_card_menu_horizontal_padding">0dp</dimen>
-    <dimen name="portrait_task_card_horz_space_big_overview">96dp</dimen>
+    <dimen name="portrait_task_card_horz_space_big_overview">132dp</dimen>
     <dimen name="portrait_modal_task_card_horz_space">60dp</dimen>
     <dimen name="landscape_task_card_horz_space">200dp</dimen>
     <dimen name="multi_window_task_card_horz_space">100dp</dimen>
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 2e018f3..d98e792 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -21,12 +21,11 @@
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
+import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.LauncherState.OVERVIEW;
 import static com.android.launcher3.LauncherState.OVERVIEW_MODAL_TASK;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.sendCustomAccessibilityEvent;
-import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.EXTENDED_CONTAINERS;
-import static com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers.ContainerCase.DEVICE_SEARCH_RESULT_CONTAINER;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_TAP;
 import static com.android.launcher3.testing.TestProtocol.HINT_STATE_ORDINAL;
 import static com.android.launcher3.testing.TestProtocol.OVERVIEW_STATE_ORDINAL;
@@ -82,7 +81,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.stream.Stream;
 
 public class QuickstepLauncher extends BaseQuickstepLauncher {
@@ -105,13 +103,11 @@
 
     @Override
     protected void logAppLaunch(ItemInfo info, InstanceId instanceId) {
-        // If the app launch is from DeviceSearchResultContainer then add the InstanceId from
-        // LiveSearchManager to recreate the AllApps search session on the server side.
-        Optional<InstanceId> logInstanceId = this.getLiveSearchManager().getLogInstanceId();
-        if (info.getContainerInfo().getContainerCase() == EXTENDED_CONTAINERS
-                && info.getContainerInfo().getExtendedContainers().getContainerCase()
-                == DEVICE_SEARCH_RESULT_CONTAINER && logInstanceId.isPresent()) {
-            instanceId = logInstanceId.get();
+        // If the app launch is from any of the surfaces in AllApps then add the InstanceId from
+        // LiveSearchManager to recreate the AllApps session on the server side.
+        if (mAllAppsSessionLogId != null && ALL_APPS.equals(
+                getStateManager().getCurrentStableState())) {
+            instanceId = mAllAppsSessionLogId;
         }
 
         StatsLogger logger = getStatsLogManager()
diff --git a/quickstep/src/com/android/quickstep/ImageActionsApi.java b/quickstep/src/com/android/quickstep/ImageActionsApi.java
index cb4d53a..8cb64c2 100644
--- a/quickstep/src/com/android/quickstep/ImageActionsApi.java
+++ b/quickstep/src/com/android/quickstep/ImageActionsApi.java
@@ -64,6 +64,20 @@
      */
     @UiThread
     public void shareWithExplicitIntent(@Nullable Rect crop, Intent intent) {
+        addImageAndSendIntent(crop, intent, false);
+    }
+
+    /**
+     * Share the image this api was constructed with using the provided intent. The implementation
+     * should set the intent's data field to the URI pointing to the image.
+     */
+    @UiThread
+    public void shareAsDataWithExplicitIntent(@Nullable Rect crop, Intent intent) {
+        addImageAndSendIntent(crop, intent, true);
+    }
+
+    @UiThread
+    private void addImageAndSendIntent(@Nullable Rect crop, Intent intent, boolean setData) {
         if (mBitmapSupplier.get() == null) {
             Log.e(TAG, "No snapshot available, not starting share.");
             return;
@@ -71,12 +85,14 @@
 
         UI_HELPER_EXECUTOR.execute(() -> persistBitmapAndStartActivity(mContext,
                 mBitmapSupplier.get(), crop, intent, (uri, intentForUri) -> {
-                    intentForUri
-                            .addFlags(FLAG_GRANT_READ_URI_PERMISSION)
-                            .putExtra(EXTRA_STREAM, uri);
+                    intentForUri.addFlags(FLAG_GRANT_READ_URI_PERMISSION);
+                    if (setData) {
+                        intentForUri.setData(uri);
+                    } else {
+                        intentForUri.putExtra(EXTRA_STREAM, uri);
+                    }
                     return new Intent[]{intentForUri};
                 }, TAG));
-
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index 844d6f5..0d2c42e 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -122,12 +122,11 @@
      */
     public static class TaskOverlay<T extends OverviewActionsView> {
 
-        private final Context mApplicationContext;
+        protected final Context mApplicationContext;
         protected final TaskThumbnailView mThumbnailView;
 
         private T mActionsView;
-        private ImageActionsApi mImageApi;
-        private boolean mIsAllowedByPolicy;
+        protected ImageActionsApi mImageApi;
 
         protected TaskOverlay(TaskThumbnailView taskThumbnailView) {
             mApplicationContext = taskThumbnailView.getContext().getApplicationContext();
@@ -153,24 +152,8 @@
 
             if (thumbnail != null) {
                 getActionsView().updateDisabledFlags(DISABLED_ROTATED, rotated);
-                final boolean isAllowedByPolicy = thumbnail.isRealSnapshot;
-
-                getActionsView().setCallbacks(new OverlayUICallbacks() {
-                    @Override
-                    public void onShare() {
-                        if (isAllowedByPolicy) {
-                            endLiveTileMode(() -> mImageApi.startShareActivity(null));
-                        } else {
-                            showBlockedByPolicyMessage();
-                        }
-                    }
-
-                    @SuppressLint("NewApi")
-                    @Override
-                    public void onScreenshot() {
-                        endLiveTileMode(() -> saveScreenshot(task));
-                    }
-                });
+                boolean isAllowedByPolicy = thumbnail.isRealSnapshot;
+                getActionsView().setCallbacks(new OverlayUICallbacksImpl(isAllowedByPolicy, task));
             }
         }
 
@@ -193,7 +176,7 @@
          * Called to save screenshot of the task thumbnail.
          */
         @SuppressLint("NewApi")
-        private void saveScreenshot(Task task) {
+        protected void saveScreenshot(Task task) {
             if (mThumbnailView.isRealSnapshot()) {
                 mImageApi.saveScreenshot(mThumbnailView.getThumbnail(),
                         getTaskSnapshotBounds(), getTaskSnapshotInsets(), task.key);
@@ -257,7 +240,7 @@
             return mThumbnailView.getScaledInsets();
         }
 
-        private void showBlockedByPolicyMessage() {
+        protected void showBlockedByPolicyMessage() {
             Toast.makeText(
                     mThumbnailView.getContext(),
                     R.string.blocked_by_policy,
@@ -279,6 +262,29 @@
                 dismissTaskMenuView(mActivity);
             }
         }
+
+        protected class OverlayUICallbacksImpl implements OverlayUICallbacks {
+            protected final boolean mIsAllowedByPolicy;
+            protected final Task mTask;
+
+            public OverlayUICallbacksImpl(boolean isAllowedByPolicy, Task task) {
+                mIsAllowedByPolicy = isAllowedByPolicy;
+                mTask = task;
+            }
+
+            public void onShare() {
+                if (mIsAllowedByPolicy) {
+                    endLiveTileMode(() -> mImageApi.startShareActivity(null));
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+
+            @SuppressLint("NewApi")
+            public void onScreenshot() {
+                endLiveTileMode(() -> saveScreenshot(mTask));
+            }
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index aafb1af..f9283a4 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -88,7 +88,7 @@
     }
 
     @Override
-    public StatsLogger logger() {
+    protected StatsLogger createLogger() {
         return new StatsCompatLogger();
     }
 
diff --git a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
index 6f16781..edce194 100644
--- a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
+++ b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
@@ -41,7 +41,6 @@
 import com.android.quickstep.BaseActivityInterface;
 import com.android.quickstep.views.RecentsView.ScrollState;
 import com.android.quickstep.views.TaskThumbnailView.PreviewPositionHelper;
-import com.android.quickstep.views.TaskView;
 import com.android.quickstep.views.TaskView.FullscreenDrawParams;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
@@ -92,7 +91,6 @@
 
     // TaskView properties
     private final FullscreenDrawParams mCurrentFullscreenParams;
-    private float mCurveScale = 1;
     public final AnimatedFloat taskPrimaryTranslation = new AnimatedFloat();
     public final AnimatedFloat taskSecondaryTranslation = new AnimatedFloat();
 
@@ -277,8 +275,6 @@
                     .getPrimaryValue(mTaskRect.left, mTaskRect.top);
             mScrollState.screenCenter = start + mScrollState.scroll + mScrollState.halfPageSize;
             mScrollState.updateInterpolation(mDp, start);
-            mCurveScale = TaskView.getCurveScaleForInterpolation(mDp,
-                    mScrollState.linearInterpolation);
         }
 
         float progress = Utilities.boundToRange(fullScreenProgress.value, 0, 1);
@@ -295,8 +291,7 @@
         mMatrix.postTranslate(insets.left, insets.top);
         mMatrix.postScale(scale, scale);
 
-        // Apply TaskView matrix: scale, translate, scroll
-        mMatrix.postScale(mCurveScale, mCurveScale, taskWidth / 2, taskHeight / 2);
+        // Apply TaskView matrix: translate, scroll
         mMatrix.postTranslate(mTaskRect.left, mTaskRect.top);
         mOrientationState.getOrientationHandler().set(mMatrix, MATRIX_POST_TRANSLATE,
                 taskPrimaryTranslation.value);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 38d488c..5d492ac 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1984,11 +1984,6 @@
             centerToOffscreenProgress = Utilities.mapRange(centerToOffscreenProgress,
                     distanceFromMidpoint / distanceToOffscreen, 1);
         }
-        // Find the task's scale based on its offscreen progress, then see how far it still needs to
-        // move to be completely offscreen.
-        Utilities.scaleRectFAboutCenter(taskPosition,
-                TaskView.getCurveScaleForInterpolation(mActivity.getDeviceProfile(),
-                        centerToOffscreenProgress));
         distanceToOffscreen = desiredLeft - taskPosition.left;
         // Finally, we need to account for RecentsView scale, because it moves tasks based on its
         // pivot. To do this, we move the task position to where it would be offscreen at scale = 1
@@ -2114,7 +2109,7 @@
             anim.play(ObjectAnimator.ofFloat(recentsView, FULLSCREEN_PROGRESS, 1));
         } else {
             // We are launching an adjacent task, so parallax the center and other adjacent task.
-            float displacementX = tv.getWidth() * (toScale - tv.getCurveScale());
+            float displacementX = tv.getWidth() * (toScale - 1f);
             float primaryTranslation = mIsRtl ? -displacementX : displacementX;
             anim.play(ObjectAnimator.ofFloat(getPageAt(centerTaskIndex),
                     mOrientationHandler.getPrimaryViewTranslate(), primaryTranslation));
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index 7e558bb..e21bf76 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -24,7 +24,6 @@
 import android.animation.ObjectAnimator;
 import android.content.Context;
 import android.graphics.Rect;
-import android.graphics.drawable.Drawable;
 import android.util.AttributeSet;
 import android.view.Gravity;
 import android.view.MotionEvent;
@@ -164,7 +163,6 @@
     }
 
     private void addMenuOptions(TaskView taskView) {
-        Drawable icon = taskView.getTask().icon.getConstantState().newDrawable();
         mTaskName.setText(TaskUtils.getTitle(getContext(), taskView.getTask()));
         mTaskName.setOnClickListener(v -> close(true));
 
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 8a8b021..a1b5533 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -44,7 +44,6 @@
 import android.animation.AnimatorListenerAdapter;
 import android.animation.AnimatorSet;
 import android.animation.ObjectAnimator;
-import android.animation.TimeInterpolator;
 import android.app.ActivityOptions;
 import android.content.Context;
 import android.content.Intent;
@@ -112,10 +111,6 @@
 
     private static final String TAG = TaskView.class.getSimpleName();
 
-    /** A curve of x from 0 to 1, where 0 is the center of the screen and 1 is the edge. */
-    private static final TimeInterpolator CURVE_INTERPOLATOR
-            = x -> (float) -Math.cos(x * Math.PI) / 2f + .5f;
-
     /**
      * The alpha of a black scrim on a page in the carousel as it leaves the screen.
      * In the resting position of the carousel, the adjacent pages have about half this scrim.
@@ -252,7 +247,6 @@
     private TaskMenuView mMenuView;
     private IconView mIconView;
     private final DigitalWellBeingToast mDigitalWellBeingToast;
-    private float mCurveScale;
     private float mFullscreenProgress;
     private float mScaleAtFullscreen = 1;
     private float mFullscreenScale = 1;
@@ -655,7 +649,7 @@
             progress = 1 - progress;
         }
         mFocusTransitionProgress = progress;
-        mSnapshotView.setDimAlphaMultipler(progress);
+        mSnapshotView.setDimAlphaMultipler(0);
         float iconScalePercentage = (float) SCALE_ICON_DURATION / DIM_ANIM_DURATION;
         float lowerClamp = invert ? 1f - iconScalePercentage : 0;
         float upperClamp = invert ? 1 : iconScalePercentage;
@@ -703,7 +697,6 @@
     }
 
     protected void resetViewTransforms() {
-        setCurveScale(1);
         // fullscreenTranslation and accumulatedTranslation should not be reset, as
         // resetViewTransforms is called during Quickswitch scrolling.
         mDismissTranslationX = mTaskOffsetTranslationX = mTaskResistanceTranslationX = 0f;
@@ -738,13 +731,6 @@
             return;
         }
 
-        float curveInterpolation =
-                CURVE_INTERPOLATOR.getInterpolation(scrollState.linearInterpolation);
-        float curveScaleForCurveInterpolation = getCurveScaleForCurveInterpolation(
-                mActivity.getDeviceProfile(), curveInterpolation);
-        mSnapshotView.setDimAlpha(curveInterpolation * MAX_PAGE_SCRIM_ALPHA);
-        setCurveScale(curveScaleForCurveInterpolation);
-
         float dwbBannerAlpha = Utilities.boundToRange(1.0f - 2 * scrollState.linearInterpolation,
                 0f, 1f);
         mDigitalWellBeingToast.updateBannerAlpha(dwbBannerAlpha);
@@ -826,20 +812,6 @@
     }
 
     /**
-     * How much to scale down pages near the edge of the screen according to linearInterpolation.
-     */
-    public static float getCurveScaleForInterpolation(DeviceProfile deviceProfile,
-            float linearInterpolation) {
-        float curveInterpolation = CURVE_INTERPOLATOR.getInterpolation(linearInterpolation);
-        return getCurveScaleForCurveInterpolation(deviceProfile, curveInterpolation);
-    }
-
-    private static float getCurveScaleForCurveInterpolation(DeviceProfile deviceProfile,
-            float curveInterpolation) {
-        return 1 - curveInterpolation * getEdgeScaleDownFactor(deviceProfile);
-    }
-
-    /**
      * How much to scale down pages near the edge of the screen.
      */
     public static float getEdgeScaleDownFactor(DeviceProfile deviceProfile) {
@@ -850,23 +822,14 @@
         }
     }
 
-    private void setCurveScale(float curveScale) {
-        mCurveScale = curveScale;
-        applyScale();
-    }
-
     private void setFullscreenScale(float fullscreenScale) {
         mFullscreenScale = fullscreenScale;
         applyScale();
     }
 
     private void applyScale() {
-        setScaleX(mCurveScale * mFullscreenScale);
-        setScaleY(mCurveScale * mFullscreenScale);
-    }
-
-    public float getCurveScale() {
-        return mCurveScale;
+        setScaleX(mFullscreenScale);
+        setScaleY(mFullscreenScale);
     }
 
     private void setDismissTranslationX(float x) {
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 785d413..0fe5432 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -216,7 +215,6 @@
         return background;
     }
 
-    @Ignore
     @Test
     @PortraitLandscape
     public void testAllAppsFromHome() throws Exception {
diff --git a/res/drawable/ic_expand_less.xml b/res/drawable/ic_expand_less.xml
new file mode 100644
index 0000000..8360cee
--- /dev/null
+++ b/res/drawable/ic_expand_less.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorHint">
+    <path
+        android:fillColor="#FF000000"
+        android:pathData="M18.59,16.41L20,15l-8,-8 -8,8 1.41,1.41L12,9.83"/>
+</vector>
diff --git a/res/drawable/ic_expand_more.xml b/res/drawable/ic_expand_more.xml
new file mode 100644
index 0000000..49e24f6
--- /dev/null
+++ b/res/drawable/ic_expand_more.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorHint">
+    <path
+        android:fillColor="#FF000000"
+        android:pathData="M5.41,7.59L4,9l8,8 8,-8 -1.41,-1.41L12,14.17"/>
+</vector>
diff --git a/res/drawable/widgets_tray_expand_button.xml b/res/drawable/widgets_tray_expand_button.xml
new file mode 100644
index 0000000..8316e0f
--- /dev/null
+++ b/res/drawable/widgets_tray_expand_button.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:state_checked="true"
+        android:drawable="@drawable/ic_expand_less" />
+    <item android:state_checked="false"
+        android:drawable="@drawable/ic_expand_more" />
+</selector>
diff --git a/res/layout/widgets_list_row_header.xml b/res/layout/widgets_list_row_header.xml
new file mode 100644
index 0000000..faff10c
--- /dev/null
+++ b/res/layout/widgets_list_row_header.xml
@@ -0,0 +1,71 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<com.android.launcher3.widget.picker.WidgetsListHeader xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:id="@+id/widgets_list_header"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:background="?android:attr/selectableItemBackground"
+    android:paddingVertical="20dp"
+    android:orientation="horizontal">
+
+    <ImageView
+        android:id="@+id/app_icon"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_marginHorizontal="16dp"
+        android:importantForAccessibility="no"
+        tools:src="@drawable/ic_corp"/>
+
+    <LinearLayout
+        android:layout_width="0dp"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_vertical"
+        android:layout_weight="1"
+        android:orientation="vertical"
+        android:focusable="true"
+        android:descendantFocusability="afterDescendants">
+
+        <TextView
+            android:id="@+id/app_title"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_vertical"
+            android:textColor="?android:attr/textColorPrimary"
+            android:textSize="16sp"
+            tools:text="App name" />
+
+        <TextView
+            android:id="@+id/app_subtitle"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            tools:text="n widgets" />
+
+    </LinearLayout>
+
+    <!-- This checkbox is not clickable. The outermost LinearLayout is responsible to handle all
+         click event and update the checkbox state. -->
+    <CheckBox
+        android:id="@+id/toggle"
+        android:layout_marginHorizontal="16dp"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_vertical"
+        android:layout_alignParentEnd="true"
+        android:clickable="false"
+        android:button="@drawable/widgets_tray_expand_button"/>
+
+</com.android.launcher3.widget.picker.WidgetsListHeader>
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index e593fb4..b19ea22 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -185,4 +185,8 @@
         <attr name="android:name" />
         <attr name="android:id" />
     </declare-styleable>
+
+    <declare-styleable name="WidgetsListRowHeader">
+        <attr name="appIconSize" format="dimension" />
+    </declare-styleable>
 </resources>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 447c9ac..c30019b 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -54,6 +54,11 @@
     <string name="add_item_request_drag_hint">Touch &amp; hold to place manually</string>
     <!-- Button label to automatically add icon on home screen [CHAR_LIMIT=50] -->
     <string name="place_automatically">Add automatically</string>
+    <!-- Label for showing the number of widgets an app has in the full widgets picker. [CHAR_LIMIT=25] -->
+    <plurals name="widgets_tray_subtitle">
+        <item quantity="one"><xliff:g id="widget_count" example="1">%1$d</xliff:g> widget</item>
+        <item quantity="other"><xliff:g id="widget_count" example="2">%1$d</xliff:g> widgets</item>
+    </plurals>
 
     <!-- All Apps -->
     <!-- Search bar text in the apps view. [CHAR_LIMIT=50] -->
diff --git a/robolectric_tests/config/robolectric.properties b/robolectric_tests/config/robolectric.properties
index 0ac997f..4e811f3 100644
--- a/robolectric_tests/config/robolectric.properties
+++ b/robolectric_tests/config/robolectric.properties
@@ -1,4 +1,4 @@
-sdk=30
+sdk=29
 shadows= \
     com.android.launcher3.shadows.LShadowAppPredictionManager \
     com.android.launcher3.shadows.LShadowAppWidgetManager \
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
new file mode 100644
index 0000000..04797a6
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -0,0 +1,269 @@
+/*
+ * 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;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.os.UserHandle;
+
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+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.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsDiffReporterTest {
+    private static final String TEST_PACKAGE_PREFIX = "com.google.test";
+    private static final WidgetListBaseRowEntryComparator COMPARATOR =
+            new WidgetListBaseRowEntryComparator();
+
+    @Mock private IconCache mIconCache;
+    @Mock private RecyclerView.Adapter mAdapter;
+
+    private InvariantDeviceProfile mTestProfile;
+    private WidgetsDiffReporter mWidgetsDiffReporter;
+    private Context mContext;
+    private WidgetsListHeaderEntry mHeaderA;
+    private WidgetsListHeaderEntry mHeaderB;
+    private WidgetsListHeaderEntry mHeaderC;
+    private WidgetsListHeaderEntry mHeaderD;
+    private WidgetsListHeaderEntry mHeaderE;
+    private WidgetsListContentEntry mContentC;
+    private WidgetsListContentEntry mContentE;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+                .getComponent().getPackageName())
+                .when(mIconCache).getTitleNoCache(any());
+
+        mContext = RuntimeEnvironment.application;
+        mWidgetsDiffReporter = new WidgetsDiffReporter(mIconCache, mAdapter);
+        mHeaderA = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A",
+                /* appName= */ "A", /* numOfWidgets= */ 3);
+        mHeaderB = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "B",
+                /* appName= */ "B", /* numOfWidgets= */ 3);
+        mHeaderC = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "C",
+                /* appName= */ "C", /* numOfWidgets= */ 3);
+        mContentC = createWidgetsContentEntry(TEST_PACKAGE_PREFIX + "C",
+                /* appName= */ "C", /* numOfWidgets= */ 3);
+        mHeaderD = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "D",
+                /* appName= */ "D", /* numOfWidgets= */ 3);
+        mHeaderE = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "E",
+                /* appName= */ "E", /* numOfWidgets= */ 3);
+        mContentE = createWidgetsContentEntry(TEST_PACKAGE_PREFIX + "E",
+                /* appName= */ "E", /* numOfWidgets= */ 3);
+    }
+
+    @Test
+    public void listNotChanged_shouldNotInvokeAnyCallbacks() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderC));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, currentList, COMPARATOR);
+
+        // THEN there is no adaptor callback.
+        verifyZeroInteractions(mAdapter);
+        // THEN the current list contains the same entries.
+        assertThat(currentList).containsExactly(mHeaderA, mHeaderB, mHeaderC);
+    }
+
+    @Test
+    public void headersOnly_emptyListToNonEmpty_shouldInvokeNotifyDataSetChanged() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>();
+
+        List<WidgetsListBaseEntry> newList = List.of(
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A", "A", 3),
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "B", "B", 3),
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "C", "C", 3));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notifyDataSetChanged is called
+        verify(mAdapter).notifyDataSetChanged();
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersOnly_nonEmptyToEmptyList_shouldInvokeNotifyDataSetChanged() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderC));
+        // GIVEN the new list is empty.
+        List<WidgetsListBaseEntry> newList = List.of();
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notifyDataSetChanged is called.
+        verify(mAdapter).notifyDataSetChanged();
+        // THEN the current list isEmpty.
+        assertThat(currentList).isEmpty();
+    }
+
+    @Test
+    public void headersOnly_itemAddedAndRemovedInTheNewList_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, D].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderD));
+        // GIVEN the new list has app headers [A, C, E].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mHeaderC, mHeaderE);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN "B" is removed from position 1.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 1);
+        // THEN "D" is removed from position 2.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 2);
+        // THEN "C" is inserted at position 1.
+        verify(mAdapter).notifyItemInserted(/* position= */ 1);
+        // THEN "E" is inserted at position 2.
+        verify(mAdapter).notifyItemInserted(/* position= */ 2);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersContentsMix_itemAddedAndRemovedInTheNewList_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has app headers [A, C content, D].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mContentC, mHeaderD);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN "B" is removed from position 1.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 1);
+        // THEN "C content" is inserted at position 1.
+        verify(mAdapter).notifyItemInserted(/* position= */ 1);
+        // THEN "D" is inserted at position 2.
+        verify(mAdapter).notifyItemInserted(/* position= */ 2);
+        // THEN "E content" is removed from position 3.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 3);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersContentsMix_userInteractWithHeader_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has app headers [A, B, E content].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mHeaderB, mContentE);
+        // GIVEN the user has interacted with B.
+        mHeaderB.setIsWidgetListShown(true);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notify "B" has been changed.
+        verify(mAdapter).notifyItemChanged(/* position= */ 1);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+
+    private WidgetsListHeaderEntry createWidgetsHeaderEntry(String packageName, String appName,
+            int numOfWidgets) {
+        List<WidgetItem> widgetItems = generateWidgetItems(packageName, numOfWidgets);
+        PackageItemInfo pInfo = createPackageItemInfo(packageName, appName,
+                widgetItems.get(0).user);
+
+        return new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems);
+    }
+
+    private WidgetsListContentEntry createWidgetsContentEntry(String packageName, String appName,
+            int numOfWidgets) {
+        List<WidgetItem> widgetItems = generateWidgetItems(packageName, numOfWidgets);
+        PackageItemInfo pInfo = createPackageItemInfo(packageName, appName,
+                widgetItems.get(0).user);
+
+        return new WidgetsListContentEntry(pInfo, /* titleSectionName= */ "", widgetItems);
+    }
+
+    private PackageItemInfo createPackageItemInfo(String packageName, String appName,
+            UserHandle userHandle) {
+        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        pInfo.title = appName;
+        pInfo.user = userHandle;
+        pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+        return pInfo;
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            WidgetItem widgetItem = new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache);
+            widgetItems.add(widgetItem);
+        }
+        return widgetItems;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
index 9bea2fb..e94b253 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -40,11 +40,13 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
@@ -56,9 +58,7 @@
 
 @RunWith(RobolectricTestRunner.class)
 public final class WidgetsListAdapterTest {
-
-    private static final String TEST_PACKAGE_1 = "com.google.test.1";
-    private static final String TEST_PACKAGE_2 = "com.google.test.2";
+    private static final String TEST_PACKAGE_PLACEHOLDER = "com.google.test";
 
     @Mock private LayoutInflater mMockLayoutInflater;
     @Mock private WidgetPreviewLoader mMockWidgetCache;
@@ -117,37 +117,76 @@
     }
 
     @Test
-    public void setWidgets_sameApp_moreWidgets_shouldNotifyItemChangedWithWidgetItemInfoDiff() {
-        // GIVEN the adapter was first populated with test package 1 & test package 2.
-        WidgetsListBaseEntry testPackage1With2WidgetsListEntry =
-                generateSampleAppWithWidgets(TEST_PACKAGE_1, /* numOfWidgets= */ 2);
-        WidgetsListBaseEntry testPackage2With2WidgetsListEntry =
-                generateSampleAppWithWidgets(TEST_PACKAGE_2, /* numOfWidgets= */ 2);
-        mAdapter.setWidgets(
-                List.of(testPackage1With2WidgetsListEntry, testPackage2With2WidgetsListEntry));
+    public void headerClick_expanded_shouldNotifyItemChange() {
+        // GIVEN a list of widgets entries:
+        // [com.google.test0, com.google.test0 content,
+        //  com.google.test1, com.google.test1 content,
+        //  com.google.test2, com.google.test2 content]
+        // The visible widgets entries: [com.google.test0, com.google.test1, com.google.test2].
+        mAdapter.setWidgets(generateSampleMap(3));
 
-        // WHEN the adapter is updated with the same list of apps but test package 2 has 3 widgets
+        // WHEN com.google.test.1 header is expanded.
+        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+
+        // THEN the visible entries list becomes:
+        // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
+        // com.google.test.1 content is inserted into position 2.
+        verify(mListener).onItemRangeInserted(eq(2), eq(1));
+    }
+
+    @Test
+    public void setWidgets_expandedApp_moreWidgets_shouldNotifyItemChangedWithWidgetItemInfoDiff() {
+        // GIVEN the adapter was first populated with com.google.test0 & com.google.test1. Each app
+        // has one widget.
+        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(2);
+        mAdapter.setWidgets(allEntries);
+        // GIVEN test com.google.test1 is expanded.
+        // Visible entries in the adapter are:
+        // [com.google.test0, com.google.test1, com.google.test1 content]
+        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+        Mockito.reset(mListener);
+
+        // WHEN the adapter is updated with the same list of apps but com.google.test1 has 2 widgets
         // now.
-        WidgetsListBaseEntry testPackage1With3WidgetsListEntry =
-                generateSampleAppWithWidgets(TEST_PACKAGE_2, /* numOfWidgets= */ 2);
-        mAdapter.setWidgets(
-                List.of(testPackage1With2WidgetsListEntry, testPackage1With3WidgetsListEntry));
+        WidgetsListContentEntry testPackage1ContentEntry =
+                (WidgetsListContentEntry) allEntries.get(3);
+        WidgetItem widgetItem = testPackage1ContentEntry.mWidgets.get(0);
+        WidgetsListContentEntry newTestPackage1ContentEntry = new WidgetsListContentEntry(
+                testPackage1ContentEntry.mPkgItem,
+                testPackage1ContentEntry.mTitleSectionName, List.of(widgetItem, widgetItem));
+        allEntries.set(3, newTestPackage1ContentEntry);
+        mAdapter.setWidgets(allEntries);
 
-        // THEN the onItemRangeChanged is invoked.
-        verify(mListener).onItemRangeChanged(eq(1), eq(1), isNull());
+        // THEN the onItemRangeChanged is invoked for "com.google.test1 content" at index 2.
+        verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
     }
 
     @Test
     public void setWidgets_hodgepodge_shouldInvokeExpectedDataObserverCallbacks() {
+        // GIVEN a widgets entry list:
+        // Index:  0|   1      | 2|      3   | 4|     5    | 6|     7    | 8|     9    |
+        //        [A, A content, B, B content, C, C content, D, D content, E, E content]
         List<WidgetsListBaseEntry> allAppsWithWidgets = generateSampleMap(5);
-        // GIVEN the current widgets list consist of [A, B, E].
+        // GIVEN the current widgets list consist of [A, A content, B, B content, E, E content].
+        // GIVEN the visible widgets list consist of [A, B, E]
         List<WidgetsListBaseEntry> currentList = List.of(
-                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1), allAppsWithWidgets.get(4));
+                // A & A content
+                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1),
+                // B & B content
+                allAppsWithWidgets.get(2), allAppsWithWidgets.get(3),
+                // E & E content
+                allAppsWithWidgets.get(8), allAppsWithWidgets.get(9));
         mAdapter.setWidgets(currentList);
 
-        // WHEN the widgets list is updated to [A, C, D].
+        // WHEN the widgets list is updated to [A, A content, C, C content, D, D content].
+        // WHEN the visible widgets list is updated to [A, C, D].
         List<WidgetsListBaseEntry> newList = List.of(
-                allAppsWithWidgets.get(0), allAppsWithWidgets.get(2), allAppsWithWidgets.get(3));
+                // A & A content
+                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1),
+                // C & C content
+                allAppsWithWidgets.get(4), allAppsWithWidgets.get(5),
+                // D & D content
+                allAppsWithWidgets.get(6), allAppsWithWidgets.get(7));
         mAdapter.setWidgets(newList);
 
         // Computation logic                           | [Intermediate list during computation]
@@ -162,15 +201,23 @@
     }
 
     /**
-     * Helper method to generate the sample widget model map that can be used for the tests
-     * @param num the number of WidgetItem the map should contain
+     * Generates a list of sample widget entries.
+     *
+     * <p>Each sample app has 1 widget only. An app is represented by 2 entries,
+     * {@link WidgetsListHeaderEntry} & {@link WidgetsListContentEntry}. Only
+     * {@link WidgetsListHeaderEntry} is always visible in the {@link WidgetsListAdapter}.
+     * {@link WidgetsListContentEntry} is only shown upon clicking the corresponding app's
+     * {@link WidgetsListHeaderEntry}. Only at most one {@link WidgetsListContentEntry} is shown at
+     * a time.
+     *
+     * @param num the number of apps that have widgets.
      */
     private ArrayList<WidgetsListBaseEntry> generateSampleMap(int num) {
         ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
         if (num <= 0) return result;
 
         for (int i = 0; i < num; i++) {
-            String packageName = "com.placeholder.apk" + i;
+            String packageName = TEST_PACKAGE_PLACEHOLDER + i;
 
             List<WidgetItem> widgetItems = generateWidgetItems(packageName, /* numOfWidgets= */ 1);
 
@@ -179,23 +226,13 @@
             pInfo.user = widgetItems.get(0).user;
             pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
+            result.add(new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems));
             result.add(new WidgetsListContentEntry(pInfo, /* titleSectionName= */ "", widgetItems));
         }
 
         return result;
     }
 
-    private WidgetsListBaseEntry generateSampleAppWithWidgets(String packageName,
-            int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
-        appInfo.title = appInfo.packageName;
-        appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
-
-        return new WidgetsListContentEntry(appInfo,
-                /* titleSectionName= */ "",
-                generateWidgetItems(packageName, numOfWidgets));
-    }
-
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
         ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
new file mode 100644
index 0000000..ae5b9a5
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -0,0 +1,173 @@
+/*
+ * 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;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.widget.FrameLayout;
+import android.widget.TextView;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.R;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListHeaderViewHolderBinder.OnHeaderClickListener;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsListHeaderViewHolderBinderTest {
+    private static final String TEST_PACKAGE = "com.google.test";
+    private static final String APP_NAME = "Test app";
+
+    private Context mContext;
+    private WidgetsListHeaderViewHolderBinder mViewHolderBinder;
+    private InvariantDeviceProfile mTestProfile;
+    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
+    // testing.
+    private ActivityController<TestActivity> mActivityController;
+    private TestActivity mTestActivity;
+    private FakeOnHeaderClickListener mFakeOnHeaderClickListener = new FakeOnHeaderClickListener();
+
+    @Mock
+    private IconCache mIconCache;
+    @Mock
+    private DeviceProfile mDeviceProfile;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        mActivityController = Robolectric.buildActivity(TestActivity.class);
+        mTestActivity = mActivityController.setup().get();
+        mTestActivity.setDeviceProfile(mDeviceProfile);
+
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
+
+        mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
+                LayoutInflater.from(mTestActivity),
+                mFakeOnHeaderClickListener);
+    }
+
+    @After
+    public void tearDown() {
+        mActivityController.destroy();
+    }
+
+    @Test
+    public void bindViewHolder_appWith3Widgets_shouldShowTheCorrectAppNameAndSubtitle() {
+        WidgetsListHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        WidgetsListHeaderEntry entry = generateSampleAppHeader(
+                APP_NAME,
+                TEST_PACKAGE,
+                /* numOfWidgets= */ 3);
+        mViewHolderBinder.bindViewHolder(viewHolder, entry);
+
+        TextView appTitle = widgetsListHeader.findViewById(R.id.app_title);
+        TextView appSubtitle = widgetsListHeader.findViewById(R.id.app_subtitle);
+        assertThat(appTitle.getText()).isEqualTo(APP_NAME);
+        assertThat(appSubtitle.getText()).isEqualTo("3 widgets");
+    }
+
+    private WidgetsListHeaderEntry generateSampleAppHeader(String appName, String packageName,
+            int numOfWidgets) {
+        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        appInfo.title = appName;
+        appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+
+        return new WidgetsListHeaderEntry(appInfo,
+                /* titleSectionName= */ "",
+                generateWidgetItems(packageName, numOfWidgets));
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            widgetItems.add(new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache));
+        }
+        return widgetItems;
+    }
+
+    private void assertWidgetCellWithLabel(View view, String label) {
+        assertThat(view).isInstanceOf(WidgetCell.class);
+        TextView widgetLabel = (TextView) view.findViewById(R.id.widget_name);
+        assertThat(widgetLabel.getText()).isEqualTo(label);
+    }
+
+    private final class FakeOnHeaderClickListener implements OnHeaderClickListener {
+
+        boolean mShowWidgets = false;
+        @Nullable  String mHeaderClickedPackage = null;
+
+        @Override
+        public void onHeaderClicked(boolean showWidgets, String packageName) {
+            mShowWidgets = showWidgets;
+            mHeaderClickedPackage = packageName;
+        }
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
index 4e9e227..ec9fde3 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
@@ -119,19 +119,6 @@
     }
 
     @Test
-    public void bindViewHolder_appWith3Widgets_shouldMatchAppTitle() {
-        WidgetsRowViewHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
-        WidgetsListContentEntry entry = generateSampleAppWithWidgets(
-                APP_NAME,
-                TEST_PACKAGE,
-                /* numOfWidgets= */ 3);
-        mViewHolderBinder.bindViewHolder(viewHolder, entry);
-
-        assertThat(viewHolder.title.getText()).isEqualTo(APP_NAME);
-    }
-
-    @Test
     public void bindViewHolder_appWith3Widgets_shouldHave3Widgets() {
         WidgetsRowViewHolder viewHolder = mViewHolderBinder.newViewHolder(
                 new FrameLayout(mTestActivity));
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index a11d0c9..86df3f8 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -50,20 +50,11 @@
 @RunWith(RobolectricTestRunner.class)
 public final class WidgetsListContentEntryTest {
     private static final String PACKAGE_NAME = "com.google.test";
-    private static final PackageItemInfo PACKAGE_ITEM_INFO = new PackageItemInfo(PACKAGE_NAME);
-    private static final ComponentName WIDGET_1 = ComponentName.createRelative(PACKAGE_NAME,
-            ".widget1");
-    private static final ComponentName WIDGET_2 = ComponentName.createRelative(PACKAGE_NAME,
-            ".widget2");
-    private static final ComponentName WIDGET_3 = ComponentName.createRelative(PACKAGE_NAME,
-            ".widget3");
-    private static final Map<ComponentName, String> WIDGETS_TO_LABELS = new HashMap();
-
-    static {
-        WIDGETS_TO_LABELS.put(WIDGET_1, "Cat");
-        WIDGETS_TO_LABELS.put(WIDGET_2, "Dog");
-        WIDGETS_TO_LABELS.put(WIDGET_3, "Bird");
-    }
+    private final PackageItemInfo mPackageItemInfo = new PackageItemInfo(PACKAGE_NAME);
+    private final ComponentName mWidget1 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget1");
+    private final ComponentName mWidget2 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget2");
+    private final ComponentName mWidget3 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget3");
+    private final Map<ComponentName, String> mWidgetsToLabels = new HashMap();
 
     @Mock private IconCache mIconCache;
 
@@ -73,6 +64,11 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+
+        mWidgetsToLabels.put(mWidget1, "Cat");
+        mWidgetsToLabels.put(mWidget2, "Dog");
+        mWidgetsToLabels.put(mWidget3, "Bird");
+
         mContext = RuntimeEnvironment.application;
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
@@ -80,7 +76,7 @@
 
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
-            return WIDGETS_TO_LABELS.get(componentWithLabel.getComponent());
+            return mWidgetsToLabels.get(componentWithLabel.getComponent());
         }).when(mIconCache).getTitleNoCache(any());
     }
 
@@ -88,14 +84,14 @@
     public void unsortedWidgets_diffLabels_shouldSortWidgetItems() {
         // GIVEN a list of widgets in unsorted order.
         // Cat 2x3
-        WidgetItem widgetItem1 = createWidgetItem(WIDGET_1, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
         // Dog 2x3
-        WidgetItem widgetItem2 = createWidgetItem(WIDGET_2, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetItem widgetItem2 = createWidgetItem(mWidget2, /* spanX= */ 2, /* spanY= */ 3);
         // Bird 2x3
-        WidgetItem widgetItem3 = createWidgetItem(WIDGET_3, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetItem widgetItem3 = createWidgetItem(mWidget3, /* spanX= */ 2, /* spanY= */ 3);
 
         // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
-        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(PACKAGE_ITEM_INFO,
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
                 /* titleSectionName= */ "T",
                 List.of(widgetItem1, widgetItem2, widgetItem3));
 
@@ -104,21 +100,21 @@
                 .containsExactly(widgetItem3, widgetItem1, widgetItem2)
                 .inOrder();
         assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
-        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(PACKAGE_ITEM_INFO);
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
     }
 
     @Test
     public void unsortedWidgets_sameLabels_differentSize_shouldSortWidgetItems() {
         // GIVEN a list of widgets in unsorted order.
         // Cat 3x3
-        WidgetItem widgetItem1 = createWidgetItem(WIDGET_1, /* spanX= */ 3, /* spanY= */ 3);
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 3, /* spanY= */ 3);
         // Cat 1x2
-        WidgetItem widgetItem2 = createWidgetItem(WIDGET_1, /* spanX= */ 1, /* spanY= */ 2);
+        WidgetItem widgetItem2 = createWidgetItem(mWidget1, /* spanX= */ 1, /* spanY= */ 2);
         // Cat 2x2
-        WidgetItem widgetItem3 = createWidgetItem(WIDGET_1, /* spanX= */ 2, /* spanY= */ 2);
+        WidgetItem widgetItem3 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 2);
 
         // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
-        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(PACKAGE_ITEM_INFO,
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
                 /* titleSectionName= */ "T",
                 List.of(widgetItem1, widgetItem2, widgetItem3));
 
@@ -128,23 +124,23 @@
                 .containsExactly(widgetItem2, widgetItem3, widgetItem1)
                 .inOrder();
         assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
-        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(PACKAGE_ITEM_INFO);
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
     }
 
     @Test
     public void unsortedWidgets_hodgepodge_shouldSortWidgetItems() {
         // GIVEN a list of widgets in unsorted order.
         // Cat 3x3
-        WidgetItem widgetItem1 = createWidgetItem(WIDGET_1, /* spanX= */ 3, /* spanY= */ 3);
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 3, /* spanY= */ 3);
         // Cat 1x2
-        WidgetItem widgetItem2 = createWidgetItem(WIDGET_1, /* spanX= */ 1, /* spanY= */ 2);
+        WidgetItem widgetItem2 = createWidgetItem(mWidget1, /* spanX= */ 1, /* spanY= */ 2);
         // Dog 2x2
-        WidgetItem widgetItem3 = createWidgetItem(WIDGET_2, /* spanX= */ 2, /* spanY= */ 2);
+        WidgetItem widgetItem3 = createWidgetItem(mWidget2, /* spanX= */ 2, /* spanY= */ 2);
         // Bird 2x2
-        WidgetItem widgetItem4 = createWidgetItem(WIDGET_3, /* spanX= */ 2, /* spanY= */ 2);
+        WidgetItem widgetItem4 = createWidgetItem(mWidget3, /* spanX= */ 2, /* spanY= */ 2);
 
         // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
-        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(PACKAGE_ITEM_INFO,
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
                 /* titleSectionName= */ "T",
                 List.of(widgetItem1, widgetItem2, widgetItem3, widgetItem4));
 
@@ -155,11 +151,11 @@
                 .containsExactly(widgetItem4, widgetItem2, widgetItem1, widgetItem3)
                 .inOrder();
         assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
-        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(PACKAGE_ITEM_INFO);
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
     }
 
     private WidgetItem createWidgetItem(ComponentName componentName, int spanX, int spanY) {
-        String label = WIDGETS_TO_LABELS.get(componentName);
+        String label = mWidgetsToLabels.get(componentName);
         ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
         widgetInfo.provider = componentName;
diff --git a/src/com/android/launcher3/BaseActivity.java b/src/com/android/launcher3/BaseActivity.java
index 5e50e27..8a03fac 100644
--- a/src/com/android/launcher3/BaseActivity.java
+++ b/src/com/android/launcher3/BaseActivity.java
@@ -141,7 +141,10 @@
         return mDeviceProfile;
     }
 
-    public final StatsLogManager getStatsLogManager() {
+    /**
+     * Returns {@link StatsLogManager} for user event logging.
+     */
+    public StatsLogManager getStatsLogManager() {
         if (mStatsLogManager == null) {
             mStatsLogManager = StatsLogManager.newInstance(this);
         }
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 21297c9..cea8cd6 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -24,7 +24,6 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
-import android.animation.ValueAnimator;
 import android.content.Context;
 import android.content.res.ColorStateList;
 import android.content.res.TypedArray;
@@ -34,8 +33,6 @@
 import android.graphics.Paint;
 import android.graphics.Path;
 import android.graphics.PointF;
-import android.graphics.PorterDuff.Mode;
-import android.graphics.PorterDuffColorFilter;
 import android.graphics.Rect;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
@@ -52,7 +49,6 @@
 
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
-import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.Launcher.OnResumeCallback;
 import com.android.launcher3.accessibility.LauncherAccessibilityDelegate;
@@ -798,7 +794,7 @@
         if (mIcon != null
                 && mIcon instanceof PlaceHolderIconDrawable
                 && iconUpdateAnimationEnabled()) {
-            animateIconUpdate((PlaceHolderIconDrawable) mIcon, icon);
+            ((PlaceHolderIconDrawable) mIcon).animateIconUpdate(icon);
         }
 
         mDisableRelayout = false;
@@ -950,28 +946,6 @@
         }
     }
 
-    private static void animateIconUpdate(PlaceHolderIconDrawable oldIcon, Drawable newIcon) {
-        int placeholderColor = oldIcon.mPaint.getColor();
-        int originalAlpha = Color.alpha(placeholderColor);
-
-        ValueAnimator iconUpdateAnimation = ValueAnimator.ofInt(originalAlpha, 0);
-        iconUpdateAnimation.setDuration(ICON_UPDATE_ANIMATION_DURATION);
-        iconUpdateAnimation.addUpdateListener(valueAnimator -> {
-            int newAlpha = (int) valueAnimator.getAnimatedValue();
-            int newColor = ColorUtils.setAlphaComponent(placeholderColor, newAlpha);
-
-            newIcon.setColorFilter(new PorterDuffColorFilter(newColor, Mode.SRC_ATOP));
-        });
-        iconUpdateAnimation.addListener(new AnimatorListenerAdapter() {
-            @Override
-            public void onAnimationEnd(Animator animation) {
-                newIcon.setColorFilter(null);
-            }
-        });
-        iconUpdateAnimation.start();
-    }
-
-
     @Override
     public void decorate(int color) {
         mHighlightColor = color;
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 2df9cbe..51e7c7d 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -40,6 +40,8 @@
 import static com.android.launcher3.dragndrop.DragLayer.ALPHA_INDEX_LAUNCHER_LOAD;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_HOME;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ENTRY;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_EXIT;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ONRESUME;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ONSTOP;
 import static com.android.launcher3.model.ItemInstallQueue.FLAG_ACTIVITY_PAUSED;
@@ -124,6 +126,8 @@
 import com.android.launcher3.keyboard.ViewGroupFocusHelper;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.FileLog;
+import com.android.launcher3.logging.InstanceId;
+import com.android.launcher3.logging.InstanceIdSequence;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.ItemInstallQueue;
@@ -352,6 +356,13 @@
 
     private SafeCloseable mUserChangedCallbackCloseable;
 
+    // New InstanceId is assigned to mAllAppsSessionLogId for each AllApps sessions.
+    // When Launcher is not in AllApps state mAllAppsSessionLogId will be null.
+    // User actions within AllApps state are logged with this InstanceId, to recreate AllApps
+    // session on the server side.
+    protected InstanceId mAllAppsSessionLogId;
+    private LauncherState mPrevLauncherState;
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         Object traceToken = TraceHelper.INSTANCE.beginSection(ON_CREATE_EVT,
@@ -1027,6 +1038,7 @@
         }
         // When multiple pages are visible, show persistent page indicator
         mWorkspace.getPageIndicator().setShouldAutoHide(!state.hasFlag(FLAG_MULTI_PAGE));
+        mPrevLauncherState = mStateManager.getCurrentStableState();
     }
 
     @Override
@@ -1050,6 +1062,17 @@
             // Clear any rotation locks when going to normal state
             getRotationHelper().setCurrentStateRequest(REQUEST_NONE);
         }
+
+        if (ALL_APPS.equals(state)) {
+            // creates new instance ID since new all apps session is started.
+            mAllAppsSessionLogId = new InstanceIdSequence().newInstanceId();
+            getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_ENTRY);
+        } else if (ALL_APPS.equals(mPrevLauncherState)
+                // Check if mLogInstanceId is not null to make sure exit event is logged only once.
+                && mAllAppsSessionLogId != null) {
+            getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_EXIT);
+            mAllAppsSessionLogId = null;
+        }
     }
 
     @Override
@@ -2818,4 +2841,9 @@
         public Configuration config;
         public Bitmap snapshot;
     }
+
+    @Override
+    public StatsLogManager getStatsLogManager() {
+        return super.getStatsLogManager().withDefaultInstanceId(mAllAppsSessionLogId);
+    }
 }
diff --git a/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java b/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
index 1d31975..6c95992 100644
--- a/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
+++ b/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
@@ -110,7 +110,6 @@
      */
     public static class SectionDecorationHandler {
         private static final int FILL_ALPHA = 0;
-        private static final int FOCUS_ALPHA = (int) (.9f * 255);
 
         protected RectF mBounds = new RectF();
         private final boolean mIsFullWidth;
@@ -123,9 +122,9 @@
 
         public SectionDecorationHandler(Context context, boolean isFullWidth) {
             mIsFullWidth = isFullWidth;
-            int endScrim = Themes.getAttrColor(context, R.attr.allAppsScrimColor);
+            int endScrim = Themes.getColorBackground(context);
             mFillcolor = ColorUtils.setAlphaComponent(endScrim, FILL_ALPHA);
-            mFocusColor = ColorUtils.setAlphaComponent(endScrim, FOCUS_ALPHA);
+            mFocusColor = endScrim;
             mRadius = Themes.getDialogCornerRadius(context);
         }
 
diff --git a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
index bb1a4c0..4876298 100644
--- a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
@@ -90,10 +90,9 @@
     public void onTabChanged(int pos) {
         super.onTabChanged(pos);
         if (mUsingTabs) {
-
             // Log tab switches only when the launcher is in AllApps state
             if (mLauncher.getStateManager().getCurrentStableState() == LauncherState.ALL_APPS) {
-                mLauncher.getLiveSearchManager().allAppsLogger()
+                mLauncher.getStatsLogManager().logger()
                         .log(pos == AdapterHolder.WORK ? LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB
                                 : LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB);
             }
@@ -109,6 +108,6 @@
     @Override
     protected void hideIme() {
         super.hideIme();
-        mLauncher.getLiveSearchManager().allAppsLogger().log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED);
+        mLauncher.getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED);
     }
 }
diff --git a/src/com/android/launcher3/allapps/search/LiveSearchManager.java b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
index 748ba50..4ef154e 100644
--- a/src/com/android/launcher3/allapps/search/LiveSearchManager.java
+++ b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
@@ -16,8 +16,6 @@
 package com.android.launcher3.allapps.search;
 
 import static com.android.launcher3.LauncherState.ALL_APPS;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ENTRY;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_EXIT;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.THREAD_POOL_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
@@ -47,9 +45,6 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.logging.InstanceId;
-import com.android.launcher3.logging.InstanceIdSequence;
-import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.SafeCloseable;
@@ -57,7 +52,6 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.Optional;
 
 /**
  * Manages Lifecycle for Live search results
@@ -74,8 +68,6 @@
     private final HashMap<ComponentKey, SearchWidgetInfoContainer> mWidgetPlaceholders =
             new HashMap<>();
     private SearchWidgetHost mSearchWidgetHost;
-    private InstanceId mLogInstanceId;
-    private LauncherState mPrevLauncherState;
 
     public LiveSearchManager(Launcher launcher) {
         mLauncher = launcher;
@@ -139,11 +131,6 @@
     }
 
     @Override
-    public void onStateTransitionStart(LauncherState toState) {
-        mPrevLauncherState = mLauncher.getStateManager().getCurrentStableState();
-    }
-
-    @Override
     public void onStateTransitionComplete(LauncherState finalState) {
         if (finalState != ALL_APPS) {
             // Clear all search session related objects
@@ -152,18 +139,6 @@
 
             clearWidgetHost();
         }
-
-        if (ALL_APPS.equals(finalState)) {
-            // creates new instance ID since new all apps session is started.
-            mLogInstanceId = new InstanceIdSequence().newInstanceId();
-            allAppsLogger().log(LAUNCHER_ALLAPPS_ENTRY);
-        } else if (ALL_APPS.equals(mPrevLauncherState)
-                // Check if mLogInstanceId is not null; to avoid NPE when LAUNCHER_ALLAPPS_EXIT is
-                // triggered multiple times
-                && mLogInstanceId != null) {
-            allAppsLogger().log(LAUNCHER_ALLAPPS_EXIT);
-            mLogInstanceId = null;
-        }
     }
 
     /**
@@ -181,14 +156,6 @@
         return () -> sliceLifeCycle.removeListener(listener);
     }
 
-    /**
-     * Returns {@link InstanceId} that should be used for logging events within search session, if
-     * available.
-     */
-    public Optional<InstanceId> getLogInstanceId() {
-        return Optional.ofNullable(mLogInstanceId);
-    }
-
     static class SearchWidgetHost extends AppWidgetHost {
         SearchWidgetHost(Context context) {
             super(context, SEARCH_APPWIDGET_HOST_ID);
@@ -316,15 +283,4 @@
         @Override
         public void onActivitySaveInstanceState(Activity activity, Bundle bundle) { }
     }
-
-    /**
-     * Returns new instance of {@link StatsLogger} pre-populated with details required to log
-     * AllApps specific user events.
-     */
-    public StatsLogger allAppsLogger() {
-        return getLogInstanceId()
-                .map(instanceId -> mLauncher.getStatsLogManager().logger()
-                        .withInstanceId(instanceId))
-                .orElse(mLauncher.getStatsLogManager().logger());
-    }
 }
diff --git a/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java b/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
index d347e8f..b6d25c4 100644
--- a/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
+++ b/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
@@ -19,10 +19,19 @@
 
 import static com.android.launcher3.graphics.IconShape.getShapePath;
 
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
+import android.animation.ValueAnimator;
 import android.content.Context;
 import android.graphics.Canvas;
+import android.graphics.Color;
 import android.graphics.Path;
+import android.graphics.PorterDuff;
+import android.graphics.PorterDuffColorFilter;
 import android.graphics.Rect;
+import android.graphics.drawable.Drawable;
+
+import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.FastBitmapDrawable;
 import com.android.launcher3.R;
@@ -53,4 +62,27 @@
         canvas.drawPath(mProgressPath, mPaint);
         canvas.restoreToCount(saveCount);
     }
+
+    /** Updates this placeholder to {@code newIcon} with animation. */
+    public void animateIconUpdate(Drawable newIcon) {
+        int placeholderColor = mPaint.getColor();
+        int originalAlpha = Color.alpha(placeholderColor);
+
+        ValueAnimator iconUpdateAnimation = ValueAnimator.ofInt(originalAlpha, 0);
+        iconUpdateAnimation.setDuration(375);
+        iconUpdateAnimation.addUpdateListener(valueAnimator -> {
+            int newAlpha = (int) valueAnimator.getAnimatedValue();
+            int newColor = ColorUtils.setAlphaComponent(placeholderColor, newAlpha);
+
+            newIcon.setColorFilter(new PorterDuffColorFilter(newColor, PorterDuff.Mode.SRC_ATOP));
+        });
+        iconUpdateAnimation.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                newIcon.setColorFilter(null);
+            }
+        });
+        iconUpdateAnimation.start();
+    }
+
 }
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index cc80a88..1266bb4 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -22,6 +22,8 @@
 
 import android.content.Context;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.R;
 import com.android.launcher3.logger.LauncherAtom.ContainerInfo;
 import com.android.launcher3.logger.LauncherAtom.FromState;
@@ -29,6 +31,8 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.ResourceBasedOverride;
 
+import java.util.Optional;
+
 /**
  * Handles the user event logging in R+.
  *
@@ -47,6 +51,7 @@
     public static final int LAUNCHER_STATE_ALLAPPS = 4;
     public static final int LAUNCHER_STATE_UNCHANGED = 5;
 
+    private InstanceId mInstanceId;
     /**
      * Returns event enum based on the two state transition information when swipe
      * gesture happens(to be removed during UserEventDispatcher cleanup).
@@ -480,16 +485,30 @@
      * Returns new logger object.
      */
     public StatsLogger logger() {
+        StatsLogger logger = createLogger();
+        Optional.ofNullable(mInstanceId).ifPresent(logger::withInstanceId);
+        return logger;
+    }
+
+    protected StatsLogger createLogger() {
         return new StatsLogger() {
         };
     }
 
     /**
+     * Sets InstanceId to every new {@link StatsLogger} object returned by {@link #logger()} when
+     * not-null.
+     */
+    public StatsLogManager withDefaultInstanceId(@Nullable InstanceId instanceId) {
+        this.mInstanceId = instanceId;
+        return this;
+    }
+
+    /**
      * Creates a new instance of {@link StatsLogManager} based on provided context.
      */
     public static StatsLogManager newInstance(Context context) {
-        StatsLogManager mgr = Overrides.getObject(StatsLogManager.class,
+        return Overrides.getObject(StatsLogManager.class,
                 context.getApplicationContext(), R.string.stats_log_manager_class);
-        return mgr;
     }
 }
diff --git a/src/com/android/launcher3/util/Themes.java b/src/com/android/launcher3/util/Themes.java
index 55d17fc..512a286 100644
--- a/src/com/android/launcher3/util/Themes.java
+++ b/src/com/android/launcher3/util/Themes.java
@@ -82,6 +82,11 @@
     }
 
     /** Returns the floating background color attribute. */
+    public static int getColorBackground(Context context) {
+        return getAttrColor(context, android.R.attr.colorBackground);
+    }
+
+    /** Returns the floating background color attribute. */
     public static int getColorBackgroundFloating(Context context) {
         return getAttrColor(context, android.R.attr.colorBackgroundFloating);
     }
diff --git a/src/com/android/launcher3/views/ScrimView.java b/src/com/android/launcher3/views/ScrimView.java
index 77cec80..c9bd284 100644
--- a/src/com/android/launcher3/views/ScrimView.java
+++ b/src/com/android/launcher3/views/ScrimView.java
@@ -42,7 +42,7 @@
  */
 public class ScrimView<T extends Launcher> extends View implements Insettable, OnChangeListener {
 
-    private static final float SCRIM_ALPHA = .75f;
+    private static final float SCRIM_ALPHA = .95f;
     protected final T mLauncher;
     private final WallpaperColorInfo mWallpaperColorInfo;
     protected final int mEndScrim;
@@ -63,6 +63,7 @@
         mWallpaperColorInfo = WallpaperColorInfo.INSTANCE.get(context);
         int endScrim = Themes.getAttrColor(context, R.attr.allAppsScrimColor);
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
+            endScrim = Themes.getColorBackgroundFloating(context);
             endScrim = ColorUtils.setAlphaComponent(endScrim, (int) (255  * SCRIM_ALPHA));
         }
         mEndScrim = endScrim;
diff --git a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
index 10ea7db..09517e1 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
@@ -16,9 +16,15 @@
 
 package com.android.launcher3.widget.model;
 
+import static java.lang.annotation.RetentionPolicy.SOURCE;
+
+import androidx.annotation.IntDef;
+
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
 
+import java.lang.annotation.Retention;
+
 /** Holder class to store the package information of an entry shown in the widgets list. */
 public abstract class WidgetsListBaseEntry {
     public final PackageItemInfo mPkgItem;
@@ -33,4 +39,22 @@
         mPkgItem = pkgItem;
         mTitleSectionName = titleSectionName;
     }
+
+    /**
+     * Returns the ranking of this entry in the
+     * {@link com.android.launcher3.widget.picker.WidgetsListAdapter}.
+     *
+     * <p>Entries with smaller value should be shown first. See
+     * {@link com.android.launcher3.widget.picker.WidgetsDiffReporter} for more details.
+     */
+    @Rank
+    public abstract int getRank();
+
+    @Retention(SOURCE)
+    @IntDef({RANK_WIDGETS_LIST_HEADER, RANK_WIDGETS_LIST_CONTENT})
+    public @interface Rank {
+    }
+
+    public static final int RANK_WIDGETS_LIST_HEADER = 1;
+    public static final int RANK_WIDGETS_LIST_CONTENT = 2;
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
index 407f194..b0cb8c7 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
@@ -41,4 +41,10 @@
     public String toString() {
         return mPkgItem.packageName + ":" + mWidgets.size();
     }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_CONTENT;
+    }
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
new file mode 100644
index 0000000..6899647
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
@@ -0,0 +1,64 @@
+/*
+ * 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.model;
+
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import java.util.Collection;
+
+/** An information holder for an app which has widgets or/and shortcuts. */
+public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
+
+    public final int widgetsCount;
+    public final int shortcutsCount;
+
+    private boolean mIsWidgetListShown = false;
+    private boolean mHasEntryUpdated = false;
+
+    public WidgetsListHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
+            Collection<WidgetItem> items) {
+        super(pkgItem, titleSectionName);
+        widgetsCount = (int) items.stream().filter(item -> item.widgetInfo != null).count();
+        shortcutsCount = Math.max(0, items.size() - widgetsCount);
+    }
+
+    /** Sets if the widgets list associated with this header is shown. */
+    public void setIsWidgetListShown(boolean isWidgetListShown) {
+        if (mIsWidgetListShown != isWidgetListShown) {
+            this.mIsWidgetListShown = isWidgetListShown;
+            mHasEntryUpdated = true;
+        } else {
+            mHasEntryUpdated = false;
+        }
+    }
+
+    /** Returns {@code true} if the widgets list associated with this header is shown. */
+    public boolean isWidgetListShown() {
+        return mIsWidgetListShown;
+    }
+
+    /** Returns {@code true} if this entry has been updated due to user interactions. */
+    public boolean hasEntryUpdated() {
+        return mHasEntryUpdated;
+    }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_HEADER;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
index 398d9ba..dbd1bdf 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
@@ -24,10 +24,12 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 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.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
 
 import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 
 /**
  * Do diff on widget's tray list items and call the {@link RecyclerView.Adapter}
@@ -50,7 +52,7 @@
      * relevant {@link androidx.recyclerview.widget.RecyclerView.RecyclerViewDataObserver} methods.
      */
     public void process(ArrayList<WidgetsListBaseEntry> currentEntries,
-            ArrayList<WidgetsListBaseEntry> newEntries,
+            List<WidgetsListBaseEntry> newEntries,
             WidgetListBaseRowEntryComparator comparator) {
         if (DEBUG) {
             Log.d(TAG, "process oldEntries#=" + currentEntries.size()
@@ -78,7 +80,7 @@
         WidgetsListBaseEntry newRowEntry = newIter.next();
 
         do {
-            int diff = comparePackageName(orgRowEntry, newRowEntry, comparator);
+            int diff = compareAppNameAndType(orgRowEntry, newRowEntry, comparator);
             if (DEBUG) {
                 Log.d(TAG, String.format("diff=%d orgRowEntry (%s) newRowEntry (%s)",
                         diff, orgRowEntry != null ? orgRowEntry.toString() : null,
@@ -106,11 +108,13 @@
                 mListener.notifyItemInserted(index);
 
             } else {
-                // same package name but,
+                // same app name & type but,
                 // did the icon, title, etc, change?
+                // or did the header view changed due to user interactions?
                 // or did the widget size and desc, span, etc change?
                 if (!isSamePackageItemInfo(orgRowEntry.mPkgItem, newRowEntry.mPkgItem)
-                        || !areWidgetsEqual(orgRowEntry, newRowEntry)) {
+                        || hasHeaderUpdated(newRowEntry)
+                        || hasWidgetsListChanged(orgRowEntry, newRowEntry)) {
                     index = currentEntries.indexOf(orgRowEntry);
                     currentEntries.set(index, newRowEntry);
                     mListener.notifyItemChanged(index);
@@ -126,10 +130,13 @@
     }
 
     /**
-     * Compare package name using the same comparator as in {@link WidgetsListAdapter}.
-     * Also handle null row pointers.
+     * Compares the app name and then entry type for the given {@link WidgetsListBaseEntry}s.
+     *
+     * @Return 0 if both entries' order is the same. Negative integer if {@code newRowEntry} should
+     *         order before {@code orgRowEntry}. Positive integer if {@code orgRowEntry} should
+     *         order before {@code newRowEntry}.
      */
-    private int comparePackageName(WidgetsListBaseEntry curRow, WidgetsListBaseEntry newRow,
+    private int compareAppNameAndType(WidgetsListBaseEntry curRow, WidgetsListBaseEntry newRow,
             WidgetListBaseRowEntryComparator comparator) {
         if (curRow == null && newRow == null) {
             throw new IllegalStateException(
@@ -141,10 +148,18 @@
         } else if (curRow != null && newRow == null) {
             return -1; // old row needs to be deleted
         }
-        return comparator.compare(curRow, newRow);
+        int diff = comparator.compare(curRow, newRow);
+        if (diff == 0) {
+            return newRow.getRank() - curRow.getRank();
+        }
+        return diff;
     }
 
-    private boolean areWidgetsEqual(WidgetsListBaseEntry curRow,
+    /**
+     * Returns {@code true} if both {@code curRow} & {@code newRow} are
+     * {@link WidgetsListContentEntry}s with a different list of widgets.
+     */
+    private boolean hasWidgetsListChanged(WidgetsListBaseEntry curRow,
             WidgetsListBaseEntry newRow) {
         if (!(curRow instanceof WidgetsListContentEntry)
                 || !(newRow instanceof WidgetsListContentEntry)) {
@@ -152,7 +167,19 @@
         }
         WidgetsListContentEntry orgRowEntry = (WidgetsListContentEntry) curRow;
         WidgetsListContentEntry newRowEntry = (WidgetsListContentEntry) newRow;
-        return orgRowEntry.mWidgets.equals(newRowEntry.mWidgets);
+        return !orgRowEntry.mWidgets.equals(newRowEntry.mWidgets);
+    }
+
+    /**
+     * Returns {@code true} if {@code newRow} is {@link WidgetsListHeaderEntry} and its content has
+     * been changed due to user interactions.
+     */
+    private boolean hasHeaderUpdated(WidgetsListBaseEntry newRow) {
+        if (!(newRow instanceof WidgetsListHeaderEntry)) {
+            return false;
+        }
+        WidgetsListHeaderEntry newRowEntry = (WidgetsListHeaderEntry) newRow;
+        return newRowEntry.hasEntryUpdated();
     }
 
     private boolean isSamePackageItemInfo(PackageItemInfo curInfo, PackageItemInfo newInfo) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 9d30842..5ec7f3b 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -24,6 +24,7 @@
 import android.view.View.OnLongClickListener;
 import android.view.ViewGroup;
 
+import androidx.annotation.Nullable;
 import androidx.recyclerview.widget.RecyclerView;
 import androidx.recyclerview.widget.RecyclerView.Adapter;
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
@@ -36,32 +37,42 @@
 import com.android.launcher3.widget.WidgetCell;
 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.picker.WidgetsListHeaderViewHolderBinder.OnHeaderClickListener;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
+import java.util.stream.Collectors;
 
 /**
- * List view adapter for the widget tray.
+ * Recycler view adapter for the widget tray.
  *
- * <p>Memory vs. Performance:
- * The less number of types of views are inserted into a {@link RecyclerView}, the more recycling
- * happens and less memory is consumed.
+ * <p>This adapter supports view binding of subclasses of {@link WidgetsListBaseEntry}. There are 2
+ * subclasses: {@link WidgetsListHeader} & {@link WidgetsListContentEntry}.
+ * {@link WidgetsListHeader} entries are always visible in the recycler view. At most one
+ * {@link WidgetsListContentEntry} is shown in the recycler view at any time. Clicking a
+ * {@link WidgetsListHeader} will result in expanding / collapsing a corresponding
+ * {@link WidgetsListContentEntry} of the same app.
  */
-public class WidgetsListAdapter extends Adapter<ViewHolder> {
+public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderClickListener {
 
     private static final String TAG = "WidgetsListAdapter";
     private static final boolean DEBUG = false;
 
     /** Uniquely identifies widgets list view type within the app. */
     private static final int VIEW_TYPE_WIDGETS_LIST = R.layout.widgets_list_row_view;
+    private static final int VIEW_TYPE_WIDGETS_HEADER = R.layout.widgets_list_row_header;
 
     private final WidgetsDiffReporter mDiffReporter;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
     private final WidgetsListRowViewHolderBinder mWidgetsListRowViewHolderBinder;
+    private final WidgetListBaseRowEntryComparator mRowComparator =
+            new WidgetListBaseRowEntryComparator();
 
-    private ArrayList<WidgetsListBaseEntry> mEntries = new ArrayList<>();
+    private List<WidgetsListBaseEntry> mAllEntries = new ArrayList<>();
+    private ArrayList<WidgetsListBaseEntry> mVisibleEntries = new ArrayList<>();
+    @Nullable private String mWidgetsContentVisiblePackage = null;
 
     public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
             WidgetPreviewLoader widgetPreviewLoader, IconCache iconCache,
@@ -70,6 +81,8 @@
         mWidgetsListRowViewHolderBinder = new WidgetsListRowViewHolderBinder(context,
                 layoutInflater, iconClickListener, iconLongClickListener, widgetPreviewLoader);
         mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListRowViewHolderBinder);
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_HEADER,
+                new WidgetsListHeaderViewHolderBinder(layoutInflater, this::onHeaderClicked));
     }
 
     /**
@@ -96,26 +109,39 @@
 
     @Override
     public int getItemCount() {
-        return mEntries.size();
+        return mVisibleEntries.size();
     }
 
     /** Gets the section name for {@link com.android.launcher3.views.RecyclerViewFastScroller}. */
     public String getSectionName(int pos) {
-        return mEntries.get(pos).mTitleSectionName;
+        return mVisibleEntries.get(pos).mTitleSectionName;
     }
 
     /** Updates the widget list. */
     public void setWidgets(List<WidgetsListBaseEntry> tempEntries) {
-        ArrayList<WidgetsListBaseEntry> newEntries = new ArrayList<>(tempEntries);
-        WidgetListBaseRowEntryComparator rowComparator = new WidgetListBaseRowEntryComparator();
-        Collections.sort(newEntries, rowComparator);
-        mDiffReporter.process(mEntries, newEntries, rowComparator);
+        mAllEntries = tempEntries.stream().sorted(mRowComparator)
+                .collect(Collectors.toList());
+        updateVisibleEntries();
+    }
+
+    private void updateVisibleEntries() {
+        mAllEntries.forEach(entry -> {
+            if (entry instanceof WidgetsListHeaderEntry) {
+                ((WidgetsListHeaderEntry) entry).setIsWidgetListShown(
+                        entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage));
+            }
+        });
+        List<WidgetsListBaseEntry> newVisibleEntries = mAllEntries.stream()
+                .filter(entry -> entry instanceof WidgetsListHeaderEntry
+                        || entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage))
+                .collect(Collectors.toList());
+        mDiffReporter.process(mVisibleEntries, newVisibleEntries, mRowComparator);
     }
 
     @Override
     public void onBindViewHolder(ViewHolder holder, int pos) {
         ViewHolderBinder viewHolderBinder = mViewHolderBinders.get(getItemViewType(pos));
-        viewHolderBinder.bindViewHolder(holder, mEntries.get(pos));
+        viewHolderBinder.bindViewHolder(holder, mVisibleEntries.get(pos));
     }
 
     @Override
@@ -148,13 +174,26 @@
 
     @Override
     public int getItemViewType(int pos) {
-        WidgetsListBaseEntry entry = mEntries.get(pos);
+        WidgetsListBaseEntry entry = mVisibleEntries.get(pos);
         if (entry instanceof WidgetsListContentEntry) {
             return VIEW_TYPE_WIDGETS_LIST;
+        } else if (entry instanceof WidgetsListHeaderEntry) {
+            return VIEW_TYPE_WIDGETS_HEADER;
         }
         throw new UnsupportedOperationException("ViewHolderBinder not found for " + entry);
     }
 
+    @Override
+    public void onHeaderClicked(boolean showWidgets, String expandedPackage) {
+        if (showWidgets) {
+            mWidgetsContentVisiblePackage = expandedPackage;
+            updateVisibleEntries();
+        } else if (expandedPackage.equals(mWidgetsContentVisiblePackage)) {
+            mWidgetsContentVisiblePackage = null;
+            updateVisibleEntries();
+        }
+    }
+
     /** Comparator for sorting WidgetListRowEntry based on package title. */
     public static class WidgetListBaseRowEntryComparator implements
             Comparator<WidgetsListBaseEntry> {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
new file mode 100644
index 0000000..823fb7b
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -0,0 +1,205 @@
+/*
+ * 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;
+
+import static com.android.launcher3.FastBitmapDrawable.newIcon;
+
+import android.content.Context;
+import android.content.res.Resources;
+import android.content.res.TypedArray;
+import android.graphics.drawable.Drawable;
+import android.util.AttributeSet;
+import android.widget.CheckBox;
+import android.widget.ImageView;
+import android.widget.LinearLayout;
+import android.widget.TextView;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.UiThread;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.FastBitmapDrawable;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.R;
+import com.android.launcher3.graphics.PlaceHolderIconDrawable;
+import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver;
+import com.android.launcher3.icons.cache.HandlerRunnable;
+import com.android.launcher3.model.data.ItemInfoWithIcon;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+/**
+ * A UI represents a header of an app shown in the full widgets tray.
+ *
+ * It is a {@link LinearLayout} which contains an app icon, an app name, a subtitle and a checkbox
+ * which indicates if the widgets content view underneath this header should be shown.
+ */
+public final class WidgetsListHeader extends LinearLayout implements ItemInfoUpdateReceiver {
+
+    private boolean mEnableIconUpdateAnimation = false;
+
+    @Nullable private HandlerRunnable mIconLoadRequest;
+    @Nullable private Drawable mIconDrawable;
+    private final int mIconSize;
+
+    private ImageView mAppIcon;
+    private TextView mTitle;
+    private TextView mSubtitle;
+
+    private CheckBox mExpandToggle;
+    private boolean mIsExpanded = false;
+
+    public WidgetsListHeader(Context context) {
+        this(context, /* attrs= */ null);
+    }
+
+    public WidgetsListHeader(Context context, @Nullable AttributeSet attrs) {
+        this(context, attrs, /* defStyle= */ 0);
+    }
+
+    public WidgetsListHeader(Context context, @Nullable AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+
+        ActivityContext activity = ActivityContext.lookupContext(context);
+        DeviceProfile grid = activity.getDeviceProfile();
+        TypedArray a = context.obtainStyledAttributes(attrs,
+                R.styleable.WidgetsListRowHeader, defStyleAttr, /* defStyleRes= */ 0);
+        mIconSize = a.getDimensionPixelSize(R.styleable.WidgetsListRowHeader_appIconSize,
+                grid.iconSizePx);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mAppIcon = findViewById(R.id.app_icon);
+        mTitle = findViewById(R.id.app_title);
+        mSubtitle = findViewById(R.id.app_subtitle);
+        mExpandToggle = findViewById(R.id.toggle);
+    }
+
+    /**
+     * Sets a {@link OnExpansionChangeListener} to get a callback when this app widgets section
+     * expands / collapses.
+     */
+    @UiThread
+    public void setOnExpandChangeListener(
+            @Nullable OnExpansionChangeListener onExpandChangeListener) {
+        // Use the entire touch area of this view to expand / collapse an app widgets section.
+        setOnClickListener(view -> {
+            setExpanded(!mIsExpanded);
+            onExpandChangeListener.onExpansionChange(mIsExpanded);
+        });
+    }
+
+    /** Sets the expand toggle to expand / collapse. */
+    @UiThread
+    public void setExpanded(boolean isExpanded) {
+        this.mIsExpanded = isExpanded;
+        mExpandToggle.setChecked(isExpanded);
+    }
+
+    /** Apply app icon, labels and tag using a generic {@link WidgetsListHeaderEntry}. */
+    @UiThread
+    public void applyFromItemInfoWithIcon(WidgetsListHeaderEntry entry) {
+        applyIconAndLabel(entry);
+    }
+
+    @UiThread
+    private void applyIconAndLabel(WidgetsListHeaderEntry entry) {
+        PackageItemInfo info = entry.mPkgItem;
+        setIcon(info);
+        setTitles(entry);
+        setExpanded(entry.isWidgetListShown());
+
+        super.setTag(info);
+
+        verifyHighRes();
+    }
+
+    private void setIcon(PackageItemInfo info) {
+        FastBitmapDrawable icon = newIcon(getContext(), info);
+        applyDrawables(icon);
+        mIconDrawable = icon;
+        if (mIconDrawable != null) {
+            mIconDrawable.setVisible(
+                    /* visible= */ getWindowVisibility() == VISIBLE && isShown(),
+                    /* restart= */ false);
+        }
+    }
+
+    private void applyDrawables(Drawable icon) {
+        icon.setBounds(0, 0, mIconSize, mIconSize);
+
+        mAppIcon.setImageDrawable(icon);
+
+        // If the current icon is a placeholder color, animate its update.
+        if (mIconDrawable != null
+                && mIconDrawable instanceof PlaceHolderIconDrawable
+                && mEnableIconUpdateAnimation) {
+            ((PlaceHolderIconDrawable) mIconDrawable).animateIconUpdate(icon);
+        }
+    }
+
+    private void setTitles(WidgetsListHeaderEntry entry) {
+        mTitle.setText(entry.mPkgItem.title);
+
+        if (entry.widgetsCount > 0) {
+            Resources resources = getContext().getResources();
+            mSubtitle.setText(resources.getQuantityString(R.plurals.widgets_tray_subtitle,
+                    entry.widgetsCount, entry.widgetsCount));
+            mSubtitle.setVisibility(VISIBLE);
+        } else {
+            mSubtitle.setVisibility(GONE);
+        }
+    }
+
+    @Override
+    public void reapplyItemInfo(ItemInfoWithIcon info) {
+        if (getTag() == info) {
+            mIconLoadRequest = null;
+            mEnableIconUpdateAnimation = true;
+
+            // Optimization: Starting in N, pre-uploads the bitmap to RenderThread.
+            info.bitmap.icon.prepareToDraw();
+
+            setIcon((PackageItemInfo) info);
+
+            mEnableIconUpdateAnimation = false;
+        }
+    }
+
+    /** Verifies that the current icon is high-res otherwise posts a request to load the icon. */
+    public void verifyHighRes() {
+        if (mIconLoadRequest != null) {
+            mIconLoadRequest.cancel();
+            mIconLoadRequest = null;
+        }
+        if (getTag() instanceof ItemInfoWithIcon) {
+            ItemInfoWithIcon info = (ItemInfoWithIcon) getTag();
+            if (info.usingLowResIcon()) {
+                mIconLoadRequest = LauncherAppState.getInstance(getContext()).getIconCache()
+                        .updateIconInBackground(this, info);
+            }
+        }
+    }
+
+    /** A listener for the widget section expansion / collapse events. */
+    public interface OnExpansionChangeListener {
+        /** Notifies that the widget section is expanded or collapsed. */
+        void onExpansionChange(boolean isExpanded);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java
new file mode 100644
index 0000000..d4e1b1c
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java
@@ -0,0 +1,32 @@
+/*
+ * 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;
+
+import androidx.recyclerview.widget.RecyclerView.ViewHolder;
+
+/**
+ * A {@link ViewHolder} for {@link WidgetsListHeader} of an app, which renders the app icon, the app
+ * name, label and a button for showing / hiding widgets.
+ */
+public final class WidgetsListHeaderHolder extends ViewHolder {
+    final WidgetsListHeader mWidgetsListHeader;
+
+    public WidgetsListHeaderHolder(WidgetsListHeader view) {
+        super(view);
+
+        mWidgetsListHeader = view;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
new file mode 100644
index 0000000..ed53e6f
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -0,0 +1,61 @@
+/*
+ * 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;
+
+import android.view.LayoutInflater;
+import android.view.ViewGroup;
+
+import com.android.launcher3.R;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+/**
+ * Binds data from {@link WidgetsListHeaderEntry} to UI elements in {@link WidgetsListHeaderHolder}.
+ */
+public final class WidgetsListHeaderViewHolderBinder implements
+        ViewHolderBinder<WidgetsListHeaderEntry, WidgetsListHeaderHolder> {
+    private final LayoutInflater mLayoutInflater;
+    private final OnHeaderClickListener mOnHeaderClickListener;
+
+    public WidgetsListHeaderViewHolderBinder(LayoutInflater layoutInflater,
+            OnHeaderClickListener onHeaderClickListener) {
+        mLayoutInflater = layoutInflater;
+        mOnHeaderClickListener = onHeaderClickListener;
+    }
+
+    @Override
+    public WidgetsListHeaderHolder newViewHolder(ViewGroup parent) {
+        WidgetsListHeader header = (WidgetsListHeader) mLayoutInflater.inflate(
+                R.layout.widgets_list_row_header, parent, false);
+
+        return new WidgetsListHeaderHolder(header);
+    }
+
+    @Override
+    public void bindViewHolder(WidgetsListHeaderHolder viewHolder, WidgetsListHeaderEntry data) {
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        widgetsListHeader.applyFromItemInfoWithIcon(data);
+        widgetsListHeader.setExpanded(data.isWidgetListShown());
+        widgetsListHeader.setOnExpandChangeListener(isExpanded ->
+                mOnHeaderClickListener.onHeaderClicked(isExpanded, data.mPkgItem.packageName));
+    }
+
+    /** A listener to be invoked when {@link WidgetsListHeader} is clicked. */
+    public interface OnHeaderClickListener {
+        /** Calls when {@link WidgetsListHeader} is clicked to show / hide widgets for a package. */
+        void onHeaderClicked(boolean showWidgets, String packageName);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
index 22a8d00..cec6b80 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
@@ -76,7 +76,7 @@
         }
 
         ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
-                R.layout.widgets_list_row_view, parent, false);
+                R.layout.widgets_scroll_container, parent, false);
 
         // if the end padding is 0, then container view (horizontal scroll view) doesn't respect
         // the end of the linear layout width + the start padding and doesn't allow scrolling.
@@ -122,9 +122,6 @@
             }
         }
 
-        // Bind the views in the application info section.
-        holder.title.applyFromItemInfoWithIcon(entry.mPkgItem);
-
         // Bind the view in the widget horizontal tray region.
         for (int i = 0; i < infoList.size(); i++) {
             WidgetCell widget = (WidgetCell) row.getChildAt(2 * i);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
index 9be079e..ae94584 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
@@ -19,20 +19,16 @@
 
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
 
-import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.R;
 
-/** A {@link ViewHolder} for a row in the full widget picker. */
+/** A {@link ViewHolder} for showing widgets of an app in the full widget picker. */
 public final class WidgetsRowViewHolder extends ViewHolder {
 
     public final ViewGroup cellContainer;
-    public final BubbleTextView title;
 
     public WidgetsRowViewHolder(ViewGroup v) {
         super(v);
 
         cellContainer = v.findViewById(R.id.widgets_cell_list);
-        title = v.findViewById(R.id.section);
-        title.setAccessibilityDelegate(null);
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index f27922b..30c9b5f 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -31,6 +31,7 @@
 import com.android.launcher3.widget.WidgetManagerHelper;
 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.picker.WidgetsDiffReporter;
 
 import java.util.ArrayList;
@@ -73,11 +74,11 @@
 
         for (Map.Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
             PackageItemInfo pkgItem = entry.getKey();
+            List<WidgetItem> widgetItems = entry.getValue();
             String sectionName = (pkgItem.title == null) ? "" :
                     indexer.computeSectionName(pkgItem.title);
-            WidgetsListContentEntry row =
-                    new WidgetsListContentEntry(pkgItem, sectionName, entry.getValue());
-            result.add(row);
+            result.add(new WidgetsListHeaderEntry(pkgItem, sectionName, widgetItems));
+            result.add(new WidgetsListContentEntry(pkgItem, sectionName, widgetItems));
         }
         return result;
     }
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index f6a4916..4c9a8e7 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -110,7 +110,6 @@
                 .launch(mDevice.getLauncherPackageName());
     }
 
-    @Ignore
     @Test
     public void testPressHomeOnAllAppsContextMenu() throws Exception {
         final AllApps allApps = mLauncher.getWorkspace().switchToAllApps();
@@ -235,7 +234,6 @@
         }
     }
 
-    @Ignore
     @Test
     @PortraitLandscape
     public void testAppIconLaunchFromAllAppsFromHome() throws Exception {
@@ -246,7 +244,6 @@
         runIconLaunchFromAllAppsTest(this, allApps);
     }
 
-    @Ignore
     @Test
     @PortraitLandscape
     public void testWidgets() throws Exception {
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 9d4ccff..737f891 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -92,9 +92,8 @@
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        widgets.
-                getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace(true, false);
+        widgets.getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .dragToWorkspace(true, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
 
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 49af616..f95abdb 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -31,6 +31,7 @@
 import com.android.launcher3.testing.TestProtocol;
 
 import java.util.Collection;
+import java.util.List;
 
 /**
  * All widgets container.
@@ -101,22 +102,28 @@
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "getting widget " + labelText + " in widgets list")) {
-            final UiObject2 widgetsContainer = verifyActiveContainer();
+            final UiObject2 fullWidgetsPicker = verifyActiveContainer();
             mLauncher.assertTrue("Widgets container didn't become scrollable",
-                    widgetsContainer.wait(Until.scrollable(true), WAIT_TIME_MS));
+                    fullWidgetsPicker.wait(Until.scrollable(true), WAIT_TIME_MS));
             final Point displaySize = mLauncher.getRealDisplaySize();
-            final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
 
+            final UiObject2 widgetsContainer = findTestAppWidgetsScrollContainer();
+            mLauncher.assertTrue("Can't locate widgets list for the test app: "
+                                    + mLauncher.getLauncherPackageName(),
+                    widgetsContainer != null);
+            final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> cells = mLauncher.getObjectsInContainer(
-                        widgetsContainer, "widgets_scroll_container");
-                mLauncher.assertTrue("Widgets doesn't have 2 rows", cells.size() >= 2);
+                final Collection<UiObject2> cells = widgetsContainer.getChildren();
+                mLauncher.assertTrue("Widgets doesn't have 2 rows: ", cells.size() >= 2);
                 for (UiObject2 cell : cells) {
                     final UiObject2 label = cell.findObject(labelSelector);
+                    // The logic below doesn't handle the case which a widget cell of the given
+                    // label is not yet visible on the horizontal scrolling container. This won't be
+                    // an issue once we get rid of the horizontal scrolling container.
                     if (label == null) continue;
 
-                    final UiObject2 widget = label.getParent().getParent();
+                    final UiObject2 widget = cell;
                     mLauncher.assertEquals(
                             "View is not WidgetCell",
                             "com.android.launcher3.widget.WidgetCell",
@@ -131,7 +138,7 @@
                             <= displaySize.y - mLauncher.getBottomGestureSize()) {
                         int visibleDelta = maxWidth - mLauncher.getVisibleBounds(widget).width();
                         if (visibleDelta > 0) {
-                            Rect parentBounds = mLauncher.getVisibleBounds(cell);
+                            Rect parentBounds = mLauncher.getVisibleBounds(cell.getParent());
                             mLauncher.linearGesture(parentBounds.centerX() + visibleDelta
                                             + mLauncher.getTouchSlop(),
                                     parentBounds.centerY(), parentBounds.centerX(),
@@ -153,4 +160,53 @@
             }
         }
     }
+
+    /** Finds the widgets list of this test app from the collapsed full widgets picker. */
+    private UiObject2 findTestAppWidgetsScrollContainer() {
+        final BySelector headerSelector = By.res(mLauncher.getLauncherPackageName(),
+                "widgets_list_header");
+        final BySelector targetAppSelector = By.clazz("android.widget.TextView").text(
+                mLauncher.getContext().getPackageName());
+        final BySelector widgetsContainerSelector = By.res(mLauncher.getLauncherPackageName(),
+                "widgets_cell_list");
+
+        boolean hasHeaderExpanded = false;
+        for (int i = 0; i < 40; i++) {
+            UiObject2 fullWidgetsPicker = verifyActiveContainer();
+
+            UiObject2 header = fullWidgetsPicker.findObject(headerSelector);
+            mLauncher.assertTrue("Can't find a widget header", header != null);
+
+            // Look for a header that has the test app name.
+            UiObject2 headerTitle = fullWidgetsPicker.findObject(targetAppSelector);
+            if (headerTitle != null) {
+                // If we find the header and it has not been expanded, let's click it to see the
+                // widgets list.
+                if (!hasHeaderExpanded) {
+                    hasHeaderExpanded = true;
+                    mLauncher.clickLauncherObject(headerTitle);
+                    // After clicking the header, the recyclerview has been updated. Let's refresh
+                    // the container UIObject2.
+                    fullWidgetsPicker = verifyActiveContainer();
+                    // Refresh headerTitle because the first instance is stale after
+                    // verifyActiveContainer call.
+                    headerTitle = fullWidgetsPicker.findObject(targetAppSelector);
+                }
+
+                // Look for a widgets list.
+                UiObject2 widgetsContainer = fullWidgetsPicker.findObject(widgetsContainerSelector);
+                if (widgetsContainer != null) {
+                    // Make sure the widgets list is fully visible on the screen.
+                    mLauncher.scrollToLastVisibleRow(fullWidgetsPicker,
+                            widgetsContainer.getChildren(), 0);
+                    return widgetsContainer;
+                }
+                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, List.of(headerTitle), 0);
+            } else {
+                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, header.getChildren(), 0);
+            }
+        }
+
+        return null;
+    }
 }