Merge "Remove unused badge logic on WidgetCell now that we don't show badges." into main
diff --git a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
index 8c4db4a..23cb8e9 100644
--- a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
+++ b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
@@ -35,23 +35,31 @@
 import android.view.WindowManager;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.launcher3.dragndrop.SimpleDragLayer;
 import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.WidgetPredictionsRequester;
 import com.android.launcher3.model.WidgetsModel;
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.popup.PopupDataProvider;
+import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.BaseWidgetSheet;
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /** An Activity that can host Launcher's widget picker. */
 public class WidgetPickerActivity extends BaseActivity {
     private static final String TAG = "WidgetPickerActivity";
-
     /**
      * Name of the extra that indicates that a widget being dragged.
      *
@@ -64,14 +72,33 @@
     // the intent, then widgets will not be filtered for size.
     private static final String EXTRA_DESIRED_WIDGET_WIDTH = "desired_widget_width";
     private static final String EXTRA_DESIRED_WIDGET_HEIGHT = "desired_widget_height";
-
+    /**
+     * Widgets currently added by the user in the UI surface.
+     * <p>This allows widget picker to exclude existing widgets from suggestions.</p>
+     */
+    private static final String EXTRA_ADDED_APP_WIDGETS = "added_app_widgets";
+    /**
+     * A unique identifier of the surface hosting the widgets;
+     * <p>"widgets" is reserved for home screen surface.</p>
+     * <p>"widgets_hub" is reserved for glanceable hub surface.</p>
+     */
+    private static final String EXTRA_UI_SURFACE = "ui_surface";
+    private static final Pattern UI_SURFACE_PATTERN =
+            Pattern.compile("^(widgets|widgets_hub)$");
     private SimpleDragLayer<WidgetPickerActivity> mDragLayer;
     private WidgetsModel mModel;
+    private LauncherAppState mApp;
+    private WidgetPredictionsRequester mWidgetPredictionsRequester;
     private final PopupDataProvider mPopupDataProvider = new PopupDataProvider(i -> {});
 
     private int mDesiredWidgetWidth;
     private int mDesiredWidgetHeight;
     private int mWidgetCategoryFilter;
+    @Nullable
+    private String mUiSurface;
+    // Widgets existing on the host surface.
+    @NonNull
+    private List<AppWidgetProviderInfo> mAddedWidgets = new ArrayList<>();
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -80,9 +107,8 @@
         getWindow().addFlags(WindowManager.LayoutParams.FLAG_SHOW_WHEN_LOCKED);
         getWindow().clearFlags(WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER);
 
-        LauncherAppState app = LauncherAppState.getInstance(this);
-        InvariantDeviceProfile idp = app.getInvariantDeviceProfile();
-
+        mApp = LauncherAppState.getInstance(this);
+        InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
         mDeviceProfile = idp.getDeviceProfile(this);
         mModel = new WidgetsModel();
 
@@ -97,6 +123,11 @@
         widgetSheet.disableNavBarScrim(true);
         widgetSheet.addOnCloseListener(this::finish);
 
+        parseIntentExtras();
+        refreshAndBindWidgets();
+    }
+
+    private void parseIntentExtras() {
         // A value of 0 for either size means that no filtering will occur in that dimension. If
         // both values are 0, then no size filtering will occur.
         mDesiredWidgetWidth =
@@ -108,7 +139,15 @@
         mWidgetCategoryFilter =
                 getIntent().getIntExtra(AppWidgetManager.EXTRA_CATEGORY_FILTER, 0);
 
-        refreshAndBindWidgets();
+        String uiSurfaceParam = getIntent().getStringExtra(EXTRA_UI_SURFACE);
+        if (uiSurfaceParam != null && UI_SURFACE_PATTERN.matcher(uiSurfaceParam).matches()) {
+            mUiSurface = uiSurfaceParam;
+        }
+        ArrayList<AppWidgetProviderInfo> addedWidgets = getIntent().getParcelableArrayListExtra(
+                EXTRA_ADDED_APP_WIDGETS, AppWidgetProviderInfo.class);
+        if (addedWidgets != null) {
+            mAddedWidgets = addedWidgets;
+        }
     }
 
     @NonNull
@@ -179,11 +218,12 @@
         };
     }
 
+    /** Updates the model with widgets and provides them after applying the provided filter. */
     private void refreshAndBindWidgets() {
         MODEL_EXECUTOR.execute(() -> {
             LauncherAppState app = LauncherAppState.getInstance(this);
             mModel.update(app, null);
-            final ArrayList<WidgetsListBaseEntry> widgets =
+            final List<WidgetsListBaseEntry> allWidgets =
                     mModel.getFilteredWidgetsListForPicker(
                             app.getContext(),
                             /*widgetItemFilter=*/ widget -> {
@@ -193,10 +233,37 @@
                                 return verdict.isAcceptable;
                             }
                     );
-            MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setAllWidgets(widgets));
+            bindWidgets(allWidgets);
+            if (mUiSurface != null) {
+                Map<PackageUserKey, List<WidgetItem>> allWidgetsMap = allWidgets.stream()
+                        .filter(WidgetsListHeaderEntry.class::isInstance)
+                        .collect(Collectors.toMap(
+                                entry -> PackageUserKey.fromPackageItemInfo(entry.mPkgItem),
+                                entry -> entry.mWidgets)
+                        );
+                mWidgetPredictionsRequester = new WidgetPredictionsRequester(app.getContext(),
+                        mUiSurface, allWidgetsMap);
+                mWidgetPredictionsRequester.request(mAddedWidgets, this::bindRecommendedWidgets);
+            }
         });
     }
 
+    private void bindWidgets(List<WidgetsListBaseEntry> widgets) {
+        MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setAllWidgets(widgets));
+    }
+
+    private void bindRecommendedWidgets(List<ItemInfo> recommendedWidgets) {
+        MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setRecommendedWidgets(recommendedWidgets));
+    }
+
+    @Override
+    protected void onDestroy() {
+        super.onDestroy();
+        if (mWidgetPredictionsRequester != null) {
+            mWidgetPredictionsRequester.clear();
+        }
+    }
+
     private WidgetAcceptabilityVerdict isWidgetAcceptable(WidgetItem widget) {
         final AppWidgetProviderInfo info = widget.widgetInfo;
         if (info == null) {
diff --git a/quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java b/quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java
new file mode 100644
index 0000000..8431396
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java
@@ -0,0 +1,233 @@
+/*
+ * Copyright (C) 2024 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.model;
+
+import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+
+import android.app.prediction.AppPredictionContext;
+import android.app.prediction.AppPredictionManager;
+import android.app.prediction.AppPredictor;
+import android.app.prediction.AppTarget;
+import android.app.prediction.AppTargetEvent;
+import android.app.prediction.AppTargetId;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.os.Bundle;
+import android.text.TextUtils;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+import androidx.annotation.WorkerThread;
+
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
+import com.android.launcher3.widget.picker.WidgetRecommendationCategoryProvider;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+/**
+ * Works with app predictor to fetch and process widget predictions displayed in a standalone
+ * widget picker activity for a UI surface.
+ */
+public class WidgetPredictionsRequester {
+    private static final int NUM_OF_RECOMMENDED_WIDGETS_PREDICATION = 20;
+    private static final String BUNDLE_KEY_ADDED_APP_WIDGETS = "added_app_widgets";
+
+    @Nullable
+    private AppPredictor mAppPredictor;
+    private final Context mContext;
+    @NonNull
+    private final String mUiSurface;
+    @NonNull
+    private final Map<PackageUserKey, List<WidgetItem>> mAllWidgets;
+
+    public WidgetPredictionsRequester(Context context, @NonNull String uiSurface,
+            @NonNull Map<PackageUserKey, List<WidgetItem>> allWidgets) {
+        mContext = context;
+        mUiSurface = uiSurface;
+        mAllWidgets = Collections.unmodifiableMap(allWidgets);
+    }
+
+    /**
+     * Requests predictions from the app predictions manager and registers the provided callback to
+     * receive updates when predictions are available.
+     *
+     * @param existingWidgets widgets that are currently added to the surface;
+     * @param callback        consumer of prediction results to be called when predictions are
+     *                        available
+     */
+    public void request(List<AppWidgetProviderInfo> existingWidgets,
+            Consumer<List<ItemInfo>> callback) {
+        Bundle bundle = buildBundleForPredictionSession(existingWidgets, mUiSurface);
+        Predicate<WidgetItem> filter = notOnUiSurfaceFilter(existingWidgets);
+
+        MODEL_EXECUTOR.execute(() -> {
+            clear();
+            AppPredictionManager apm = mContext.getSystemService(AppPredictionManager.class);
+            if (apm == null) {
+                return;
+            }
+
+            mAppPredictor = apm.createAppPredictionSession(
+                    new AppPredictionContext.Builder(mContext)
+                            .setUiSurface(mUiSurface)
+                            .setExtras(bundle)
+                            .setPredictedTargetCount(NUM_OF_RECOMMENDED_WIDGETS_PREDICATION)
+                            .build());
+            mAppPredictor.registerPredictionUpdates(MODEL_EXECUTOR,
+                    targets -> bindPredictions(targets, filter, callback));
+            mAppPredictor.requestPredictionUpdate();
+        });
+    }
+
+    /**
+     * Returns a bundle that can be passed in a prediction session
+     *
+     * @param addedWidgets widgets that are already added by the user in the ui surface
+     * @param uiSurface    a unique identifier of the surface hosting widgets; format
+     *                     "widgets_xx"; note - "widgets" is reserved for home screen surface.
+     */
+    @VisibleForTesting
+    static Bundle buildBundleForPredictionSession(List<AppWidgetProviderInfo> addedWidgets,
+            String uiSurface) {
+        Bundle bundle = new Bundle();
+        ArrayList<AppTargetEvent> addedAppTargetEvents = new ArrayList<>();
+        for (AppWidgetProviderInfo info : addedWidgets) {
+            ComponentName componentName = info.provider;
+            AppTargetEvent appTargetEvent = buildAppTargetEvent(uiSurface, info, componentName);
+            addedAppTargetEvents.add(appTargetEvent);
+        }
+        bundle.putParcelableArrayList(BUNDLE_KEY_ADDED_APP_WIDGETS, addedAppTargetEvents);
+        return bundle;
+    }
+
+    /**
+     * Builds the AppTargetEvent for added widgets in a form that can be passed to the widget
+     * predictor.
+     * Also see {@link PredictionHelper}
+     */
+    private static AppTargetEvent buildAppTargetEvent(String uiSurface, AppWidgetProviderInfo info,
+            ComponentName componentName) {
+        AppTargetId appTargetId = new AppTargetId("widget:" + componentName.getPackageName());
+        AppTarget appTarget = new AppTarget.Builder(appTargetId, componentName.getPackageName(),
+                /*user=*/ info.getProfile()).setClassName(componentName.getClassName()).build();
+        return new AppTargetEvent.Builder(appTarget, AppTargetEvent.ACTION_PIN)
+                .setLaunchLocation(uiSurface).build();
+    }
+
+    /**
+     * Returns a filter to match {@link WidgetItem}s that don't exist on the UI surface.
+     */
+    @NonNull
+    @VisibleForTesting
+    static Predicate<WidgetItem> notOnUiSurfaceFilter(
+            List<AppWidgetProviderInfo> existingWidgets) {
+        Set<ComponentKey> existingComponentKeys = existingWidgets.stream().map(
+                widget -> new ComponentKey(widget.provider, widget.getProfile())).collect(
+                Collectors.toSet());
+        return widgetItem -> !existingComponentKeys.contains(widgetItem);
+    }
+
+    /** Provides the predictions returned by the predictor to the registered callback. */
+    @WorkerThread
+    private void bindPredictions(List<AppTarget> targets, Predicate<WidgetItem> filter,
+            Consumer<List<ItemInfo>> callback) {
+        List<WidgetItem> filteredPredictions = filterPredictions(targets, mAllWidgets, filter);
+        List<ItemInfo> mappedPredictions = mapWidgetItemsToItemInfo(filteredPredictions);
+
+        MAIN_EXECUTOR.execute(() -> callback.accept(mappedPredictions));
+    }
+
+    /**
+     * Applies the provided filter (e.g. widgets not on workspace) on the predictions returned by
+     * the predictor.
+     */
+    @VisibleForTesting
+    static List<WidgetItem> filterPredictions(List<AppTarget> predictions,
+            Map<PackageUserKey, List<WidgetItem>> allWidgets, Predicate<WidgetItem> filter) {
+        List<WidgetItem> servicePredictedItems = new ArrayList<>();
+        List<WidgetItem> localFilteredWidgets = new ArrayList<>();
+
+        for (AppTarget prediction : predictions) {
+            List<WidgetItem> widgetsInPackage = allWidgets.get(
+                    new PackageUserKey(prediction.getPackageName(), prediction.getUser()));
+            if (widgetsInPackage == null || widgetsInPackage.isEmpty()) {
+                continue;
+            }
+            String className = prediction.getClassName();
+            if (!TextUtils.isEmpty(className)) {
+                WidgetItem item = widgetsInPackage.stream()
+                        .filter(w -> className.equals(w.componentName.getClassName()))
+                        .filter(filter)
+                        .findFirst().orElse(null);
+                if (item != null) {
+                    servicePredictedItems.add(item);
+                    continue;
+                }
+            }
+            // No widget was added by the service, try local filtering
+            widgetsInPackage.stream().filter(filter).findFirst()
+                    .ifPresent(localFilteredWidgets::add);
+        }
+        if (servicePredictedItems.isEmpty()) {
+            servicePredictedItems.addAll(localFilteredWidgets);
+        }
+
+        return servicePredictedItems;
+    }
+
+    /**
+     * Converts the list of {@link WidgetItem}s to the list of {@link ItemInfo}s.
+     */
+    private List<ItemInfo> mapWidgetItemsToItemInfo(List<WidgetItem> widgetItems) {
+        List<ItemInfo> items;
+        if (enableCategorizedWidgetSuggestions()) {
+            WidgetRecommendationCategoryProvider categoryProvider =
+                    WidgetRecommendationCategoryProvider.newInstance(mContext);
+            items = widgetItems.stream()
+                    .map(it -> new PendingAddWidgetInfo(it.widgetInfo, CONTAINER_WIDGETS_PREDICTION,
+                            categoryProvider.getWidgetRecommendationCategory(mContext, it)))
+                    .collect(Collectors.toList());
+        } else {
+            items = widgetItems.stream().map(it -> new PendingAddWidgetInfo(it.widgetInfo,
+                    CONTAINER_WIDGETS_PREDICTION)).collect(Collectors.toList());
+        }
+        return items;
+    }
+
+    /** Cleans up any open prediction sessions. */
+    public void clear() {
+        if (mAppPredictor != null) {
+            mAppPredictor.destroy();
+            mAppPredictor = null;
+        }
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
index f9a8c99..176091a 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
@@ -19,7 +19,6 @@
 
 import static com.android.launcher3.LauncherState.BACKGROUND_APP;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.os.Debug;
 import android.os.SystemProperties;
@@ -136,9 +135,6 @@
             Log.d(TAG, "setVisibleFreeformTasksCount: visibleTasksCount=" + visibleTasksCount
                     + " currentValue=" + mVisibleFreeformTasksCount);
         }
-        if (!enableDesktopWindowingMode()) {
-            return;
-        }
 
         if (visibleTasksCount != mVisibleFreeformTasksCount) {
             final boolean wasVisible = mVisibleFreeformTasksCount > 0;
@@ -180,9 +176,6 @@
             Log.d(TAG, "setOverviewStateEnabled: enabled=" + overviewStateEnabled
                     + " currentValue=" + mInOverviewState);
         }
-        if (!enableDesktopWindowingMode()) {
-            return;
-        }
         if (overviewStateEnabled != mInOverviewState) {
             mInOverviewState = overviewStateEnabled;
             if (mInOverviewState) {
@@ -202,9 +195,6 @@
             Log.d(TAG, "setBackgroundStateEnabled: enabled=" + backgroundStateEnabled
                     + " currentValue=" + mBackgroundStateEnabled);
         }
-        if (!enableDesktopWindowingMode()) {
-            return;
-        }
         if (backgroundStateEnabled != mBackgroundStateEnabled) {
             mBackgroundStateEnabled = backgroundStateEnabled;
             if (mBackgroundStateEnabled) {
@@ -229,9 +219,6 @@
      * Notify controller that recents gesture has started.
      */
     public void setRecentsGestureStart() {
-        if (!enableDesktopWindowingMode()) {
-            return;
-        }
         if (DEBUG) {
             Log.d(TAG, "setRecentsGestureStart");
         }
@@ -243,9 +230,6 @@
      * {@link com.android.quickstep.GestureState.GestureEndTarget}
      */
     public void setRecentsGestureEnd(@Nullable GestureState.GestureEndTarget endTarget) {
-        if (!enableDesktopWindowingMode()) {
-            return;
-        }
         if (DEBUG) {
             Log.d(TAG, "setRecentsGestureEnd: endTarget=" + endTarget);
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index bed85d7..d89f49b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -15,8 +15,6 @@
  */
 package com.android.launcher3.taskbar;
 
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
-
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
 
@@ -117,9 +115,7 @@
         DesktopVisibilityController desktopController =
                 LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
         final boolean onDesktop =
-                enableDesktopWindowingMode()
-                        && desktopController != null
-                        && desktopController.areFreeformTasksVisible();
+                desktopController != null && desktopController.areFreeformTasksVisible();
 
         if (mModel.isTaskListValid(mTaskListChangeId)) {
             // When we are opening the KQS with no focus override, check if the first task is
@@ -158,14 +154,12 @@
 
         // Hide all desktop tasks and show them on the hidden tile
         int hiddenDesktopTasks = 0;
-        if (enableDesktopWindowingMode()) {
-            DesktopTask desktopTask = findDesktopTask(tasks);
-            if (desktopTask != null) {
-                hiddenDesktopTasks = desktopTask.tasks.size();
-                tasks = tasks.stream()
-                        .filter(t -> !(t instanceof DesktopTask))
-                        .collect(Collectors.toCollection(ArrayList<GroupTask>::new));
-            }
+        DesktopTask desktopTask = findDesktopTask(tasks);
+        if (desktopTask != null) {
+            hiddenDesktopTasks = desktopTask.tasks.size();
+            tasks = tasks.stream()
+                    .filter(t -> !(t instanceof DesktopTask))
+                    .collect(Collectors.toCollection(ArrayList<GroupTask>::new));
         }
         mTasks = tasks.stream()
                 .limit(MAX_TASKS)
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index def6287..8c8dd1f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -21,7 +21,6 @@
 import static com.android.launcher3.taskbar.TaskbarEduTooltipControllerKt.TOOLTIP_STEP_FEATURES;
 import static com.android.launcher3.taskbar.TaskbarLauncherStateController.FLAG_VISIBLE;
 import static com.android.quickstep.TaskAnimationManager.ENABLE_SHELL_TRANSITIONS;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.animation.Animator;
 import android.animation.AnimatorSet;
@@ -220,9 +219,7 @@
         DesktopVisibilityController desktopController =
                 LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
         final boolean onDesktop =
-                enableDesktopWindowingMode()
-                        && desktopController != null
-                        && desktopController.areFreeformTasksVisible();
+                desktopController != null && desktopController.areFreeformTasksVisible();
         if (onDesktop) {
             isVisible = false;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
index e293ad4..03f55ca 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
@@ -27,7 +27,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_OVERVIEW_BUTTON_LONGPRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_OVERVIEW_BUTTON_TAP;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_HOME_KEY;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
@@ -281,12 +280,10 @@
     private void navigateHome() {
         TaskUtils.closeSystemWindowsAsync(CLOSE_SYSTEM_WINDOWS_REASON_HOME_KEY);
 
-        if (enableDesktopWindowingMode()) {
-            DesktopVisibilityController desktopVisibilityController =
-                    LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
-            if (desktopVisibilityController != null) {
-                desktopVisibilityController.onHomeActionTriggered();
-            }
+        DesktopVisibilityController desktopVisibilityController =
+                LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
+        if (desktopVisibilityController != null) {
+            desktopVisibilityController.onHomeActionTriggered();
         }
 
         mService.getOverviewCommandHelper().addCommand(OverviewCommandHelper.TYPE_HOME);
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 8b923ad..c2d3d2c 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -216,7 +216,7 @@
     private FixedContainerItems mAllAppsPredictions;
     private HotseatPredictionController mHotseatPredictionController;
     private DepthController mDepthController;
-    private DesktopVisibilityController mDesktopVisibilityController;
+    private @Nullable DesktopVisibilityController mDesktopVisibilityController;
     private QuickstepTransitionManager mAppTransitionManager;
     private OverviewActionsView mActionsView;
     private TISBindHelper mTISBindHelper;
@@ -276,8 +276,8 @@
 
         mTISBindHelper = new TISBindHelper(this, this::onTISConnected);
         mDepthController = new DepthController(this);
-        mDesktopVisibilityController = new DesktopVisibilityController(this);
         if (enableDesktopWindowingMode()) {
+            mDesktopVisibilityController = new DesktopVisibilityController(this);
             mDesktopVisibilityController.registerSystemUiListener();
             mSplitSelectStateController.initSplitFromDesktopController(this);
         }
@@ -936,15 +936,13 @@
 
     @Override
     public void setResumed() {
-        if (enableDesktopWindowingMode()) {
-            DesktopVisibilityController controller = mDesktopVisibilityController;
-            if (controller != null && controller.areFreeformTasksVisible()
-                    && !controller.isRecentsGestureInProgress()) {
-                // Return early to skip setting activity to appear as resumed
-                // TODO(b/255649902): shouldn't be needed when we have a separate launcher state
-                //  for desktop that we can use to control other parts of launcher
-                return;
-            }
+        if (mDesktopVisibilityController != null
+                && mDesktopVisibilityController.areFreeformTasksVisible()
+                && !mDesktopVisibilityController.isRecentsGestureInProgress()) {
+            // Return early to skip setting activity to appear as resumed
+            // TODO(b/255649902): shouldn't be needed when we have a separate launcher state
+            //  for desktop that we can use to control other parts of launcher
+            return;
         }
         super.setResumed();
     }
@@ -1078,6 +1076,7 @@
         return mDepthController;
     }
 
+    @Nullable
     public DesktopVisibilityController getDesktopVisibilityController() {
         return mDesktopVisibilityController;
     }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java b/quickstep/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java
index a443c00..547de77 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/states/BackgroundAppState.java
@@ -18,7 +18,6 @@
 import static com.android.launcher3.Flags.enableScalingRevealHomeAnimation;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
 import static com.android.quickstep.TaskAnimationManager.ENABLE_SHELL_TRANSITIONS;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.content.Context;
 import android.graphics.Color;
@@ -92,8 +91,7 @@
 
     @Override
     protected float getDepthUnchecked(Context context) {
-        if (enableDesktopWindowingMode()
-                && Launcher.getLauncher(context).areFreeformTasksVisible()) {
+        if (Launcher.getLauncher(context).areFreeformTasksVisible()) {
             // Don't blur the background while freeform tasks are visible
             return BaseDepthController.DEPTH_0_PERCENT;
         } else if (enableScalingRevealHomeAnimation()) {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/states/QuickSwitchState.java b/quickstep/src/com/android/launcher3/uioverrides/states/QuickSwitchState.java
index 2587395..7fb811d 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/states/QuickSwitchState.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/states/QuickSwitchState.java
@@ -16,7 +16,6 @@
 package com.android.launcher3.uioverrides.states;
 
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.graphics.Color;
 
@@ -46,11 +45,9 @@
 
     @Override
     public int getWorkspaceScrimColor(Launcher launcher) {
-        if (enableDesktopWindowingMode()) {
-            if (launcher.areFreeformTasksVisible()) {
-                // No scrim while freeform tasks are visible
-                return Color.TRANSPARENT;
-            }
+        if (launcher.areFreeformTasksVisible()) {
+            // No scrim while freeform tasks are visible
+            return Color.TRANSPARENT;
         }
         DeviceProfile dp = launcher.getDeviceProfile();
         if (dp.isTaskbarPresentInApps) {
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index dc1c6a6..62e823a 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -61,7 +61,6 @@
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.ON_SETTLED_ON_END_TARGET;
 import static com.android.quickstep.views.RecentsView.UPDATE_SYSUI_FLAGS_THRESHOLD;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -947,7 +946,7 @@
     public void onRecentsAnimationStart(RecentsAnimationController controller,
             RecentsAnimationTargets targets) {
         super.onRecentsAnimationStart(controller, targets);
-        if (enableDesktopWindowingMode() && targets.hasDesktopTasks()) {
+        if (targets.hasDesktopTasks()) {
             mRemoteTargetHandles = mTargetGluer.assignTargetsForDesktop(targets);
         } else {
             int untrimmedAppCount = mRemoteTargetHandles.length;
@@ -1170,13 +1169,11 @@
                 mStateCallback.setState(STATE_SCALED_CONTROLLER_HOME | STATE_CAPTURE_SCREENSHOT);
                 // Notify the SysUI to use fade-in animation when entering PiP
                 SystemUiProxy.INSTANCE.get(mContext).setPipAnimationTypeToAlpha();
-                if (enableDesktopWindowingMode()) {
+                DesktopVisibilityController desktopVisibilityController =
+                        mActivityInterface.getDesktopVisibilityController();
+                if (desktopVisibilityController != null) {
                     // Notify the SysUI to stash desktop apps if they are visible
-                    DesktopVisibilityController desktopVisibilityController =
-                            mActivityInterface.getDesktopVisibilityController();
-                    if (desktopVisibilityController != null) {
-                        desktopVisibilityController.onHomeActionTriggered();
-                    }
+                    desktopVisibilityController.onHomeActionTriggered();
                 }
                 break;
             case RECENTS:
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 24c99e3..a3f6be0 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -25,7 +25,6 @@
 import static com.android.quickstep.GestureState.GestureEndTarget.RECENTS;
 import static com.android.quickstep.util.RecentsAtomicAnimationFactory.INDEX_RECENTS_FADE_ANIM;
 import static com.android.quickstep.util.RecentsAtomicAnimationFactory.INDEX_RECENTS_TRANSLATE_X_ANIM;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 import static com.android.quickstep.views.RecentsView.ADJACENT_PAGE_HORIZONTAL_OFFSET;
 import static com.android.quickstep.views.RecentsView.FULLSCREEN_PROGRESS;
 import static com.android.quickstep.views.RecentsView.RECENTS_SCALE_PROPERTY;
@@ -109,19 +108,17 @@
         if (endTarget != null) {
             // We were on our way to this state when we got canceled, end there instead.
             startState = stateFromGestureEndTarget(endTarget);
-            if (enableDesktopWindowingMode()) {
-                DesktopVisibilityController controller = getDesktopVisibilityController();
-                if (controller != null && controller.areFreeformTasksVisible()
-                        && endTarget == LAST_TASK) {
-                    // When we are cancelling the transition and going back to last task, move to
-                    // rest state instead when desktop tasks are visible.
-                    // If a fullscreen task is visible, launcher goes to normal state when the
-                    // activity is stopped. This does not happen when freeform tasks are visible
-                    // on top of launcher. Force the launcher state to rest state here.
-                    startState = activity.getStateManager().getRestState();
-                    // Do not animate the transition
-                    activityVisible = false;
-                }
+            DesktopVisibilityController controller = getDesktopVisibilityController();
+            if (controller != null && controller.areFreeformTasksVisible()
+                    && endTarget == LAST_TASK) {
+                // When we are cancelling the transition and going back to last task, move to
+                // rest state instead when desktop tasks are visible.
+                // If a fullscreen task is visible, launcher goes to normal state when the
+                // activity is stopped. This does not happen when freeform tasks are visible
+                // on top of launcher. Force the launcher state to rest state here.
+                startState = activity.getStateManager().getRestState();
+                // Do not animate the transition
+                activityVisible = false;
             }
         }
         activity.getStateManager().goToState(startState, activityVisible);
diff --git a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
index fa425f1..79b09fd 100644
--- a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
+++ b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
@@ -100,7 +100,7 @@
     private final int mWindowScaleMarginX;
     private float mWindowScaleEndCornerRadius;
     private float mWindowScaleStartCornerRadius;
-    private final Interpolator mProgressInterpolator = Interpolators.STANDARD_DECELERATE;
+    private final Interpolator mProgressInterpolator = Interpolators.BACK_GESTURE;
     private final Interpolator mVerticalMoveInterpolator = new DecelerateInterpolator();
     private final PointF mInitialTouchPos = new PointF();
 
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 4c1b1e0..879fccb 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -270,9 +270,13 @@
 
         int numVisibleTasks = 0;
         for (GroupedRecentTaskInfo rawTask : rawTasks) {
-            if (enableDesktopWindowingMode() && rawTask.getType() == TYPE_FREEFORM) {
-                GroupTask desktopTask = createDesktopTask(rawTask);
-                allTasks.add(desktopTask);
+            if (rawTask.getType() == TYPE_FREEFORM) {
+                // TYPE_FREEFORM tasks is only created when enableDesktopWindowingMode() is true,
+                // leftover TYPE_FREEFORM tasks created when flag was on should be ignored.
+                if (enableDesktopWindowingMode()) {
+                    GroupTask desktopTask = createDesktopTask(rawTask);
+                    allTasks.add(desktopTask);
+                }
                 continue;
             }
             ActivityManager.RecentTaskInfo taskInfo1 = rawTask.getTaskInfo1();
diff --git a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
index ffbb064..0ce4d0a 100644
--- a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
+++ b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
@@ -17,7 +17,6 @@
 package com.android.quickstep;
 
 import static com.android.quickstep.util.SplitScreenUtils.convertShellSplitBoundsToLauncher;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 import static com.android.wm.shell.util.SplitBounds.KEY_EXTRA_SPLIT_BOUNDS;
 
 import android.app.WindowConfiguration;
@@ -68,17 +67,15 @@
      * running tasks
      */
     public RemoteTargetGluer(Context context, BaseActivityInterface sizingStrategy) {
-        if (enableDesktopWindowingMode()) {
-            DesktopVisibilityController desktopVisibilityController =
-                    LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
-            if (desktopVisibilityController != null) {
-                int visibleTasksCount = desktopVisibilityController.getVisibleFreeformTasksCount();
-                if (visibleTasksCount > 0) {
-                    // Allocate +1 to account for a new task added to the desktop mode
-                    int numHandles = visibleTasksCount + 1;
-                    init(context, sizingStrategy, numHandles, true /* forDesktop */);
-                    return;
-                }
+        DesktopVisibilityController desktopVisibilityController =
+                LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
+        if (desktopVisibilityController != null) {
+            int visibleTasksCount = desktopVisibilityController.getVisibleFreeformTasksCount();
+            if (visibleTasksCount > 0) {
+                // Allocate +1 to account for a new task added to the desktop mode
+                int numHandles = visibleTasksCount + 1;
+                init(context, sizingStrategy, numHandles, true /* forDesktop */);
+                return;
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index 4089498..bff5a25 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -35,7 +35,6 @@
 import static com.android.quickstep.util.SplitSelectDataHolder.SPLIT_TASK_PENDINGINTENT;
 import static com.android.quickstep.util.SplitSelectDataHolder.SPLIT_TASK_SHORTCUT;
 import static com.android.quickstep.util.SplitSelectDataHolder.SPLIT_TASK_TASK;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 import static com.android.wm.shell.common.split.SplitScreenConstants.KEY_EXTRA_WIDGET_INTENT;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SNAP_TO_50_50;
 
@@ -978,7 +977,6 @@
                 @Override
                 public boolean onRequestSplitSelect(ActivityManager.RunningTaskInfo taskInfo,
                         int splitPosition, Rect taskBounds) {
-                    if (!enableDesktopWindowingMode()) return false;
                     MAIN_EXECUTOR.execute(() -> enterSplitSelect(taskInfo, splitPosition,
                             taskBounds));
                     return true;
diff --git a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
index 2282c46..16d707b 100644
--- a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
@@ -17,7 +17,6 @@
 package com.android.quickstep.util;
 
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -197,8 +196,6 @@
     }
 
     private boolean shouldIgnoreSecondSplitLaunch() {
-        return (!FeatureFlags.enableSplitContextually()
-                && !enableDesktopWindowingMode())
-                || !mController.isSplitSelectActive();
+        return !FeatureFlags.enableSplitContextually() || !mController.isSplitSelectActive();
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index cd4fab6..0a3d2a0 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -26,7 +26,6 @@
 import static com.android.launcher3.LauncherState.OVERVIEW_SPLIT_SELECT;
 import static com.android.launcher3.LauncherState.SPRING_LOADED;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SPLIT_SELECTION_EXIT_HOME;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.annotation.TargetApi;
 import android.content.Context;
@@ -265,11 +264,11 @@
 
     @Override
     public void onGestureAnimationEnd() {
-        DesktopVisibilityController desktopVisibilityController = null;
+        DesktopVisibilityController desktopVisibilityController =
+                mActivity.getDesktopVisibilityController();
         boolean showDesktopApps = false;
         GestureState.GestureEndTarget endTarget = null;
-        if (enableDesktopWindowingMode()) {
-            desktopVisibilityController = mActivity.getDesktopVisibilityController();
+        if (desktopVisibilityController != null) {
             endTarget = mCurrentGestureEndTarget;
             if (endTarget == GestureState.GestureEndTarget.LAST_TASK
                     && desktopVisibilityController.areFreeformTasksVisible()) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 776abd7..2fcca87 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -3965,10 +3965,8 @@
         // tasks.
         mActionsView.updateForGroupedTask(isCurrentSplit);
 
-        if (enableDesktopWindowingMode()) {
-            boolean isCurrentDesktop = getCurrentPageTaskView() instanceof DesktopTaskView;
-            mActionsView.updateHiddenFlags(HIDDEN_DESKTOP, isCurrentDesktop);
-        }
+        boolean isCurrentDesktop = getCurrentPageTaskView() instanceof DesktopTaskView;
+        mActionsView.updateHiddenFlags(HIDDEN_DESKTOP, isCurrentDesktop);
     }
 
     /** Returns if app pairs are supported in this launcher. Overridden in subclasses. */
@@ -4690,9 +4688,7 @@
         mSplitSelectStateController.setAnimateCurrentTaskDismissal(
                 true /*animateCurrentTaskDismissal*/);
         mSplitHiddenTaskViewIndex = indexOfChild(taskView);
-        if (enableDesktopWindowingMode()) {
-            updateDesktopTaskVisibility(false /* visible */);
-        }
+        updateDesktopTaskVisibility(false /* visible */);
     }
 
     /**
@@ -4714,9 +4710,7 @@
         mSplitSelectStateController.setInitialTaskSelect(splitSelectSource.intent,
                 splitSelectSource.position.stagePosition, splitSelectSource.itemInfo,
                 splitSelectSource.splitEvent, splitSelectSource.alreadyRunningTaskId);
-        if (enableDesktopWindowingMode()) {
-            updateDesktopTaskVisibility(false /* visible */);
-        }
+        updateDesktopTaskVisibility(false /* visible */);
     }
 
     private void updateDesktopTaskVisibility(boolean visible) {
@@ -4920,9 +4914,7 @@
             mSplitHiddenTaskView.setThumbnailVisibility(VISIBLE, INVALID_TASK_ID);
             mSplitHiddenTaskView = null;
         }
-        if (enableDesktopWindowingMode()) {
-            updateDesktopTaskVisibility(true /* visible */);
-        }
+        updateDesktopTaskVisibility(true /* visible */);
     }
 
     private void safeRemoveDragLayerView(@Nullable View viewToRemove) {
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt
new file mode 100644
index 0000000..5c7b4ab
--- /dev/null
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt
@@ -0,0 +1,221 @@
+/*
+ * Copyright (C) 2024 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.model
+
+import android.app.prediction.AppTarget
+import android.app.prediction.AppTargetEvent
+import android.app.prediction.AppTargetId
+import android.appwidget.AppWidgetProviderInfo
+import android.content.ComponentName
+import android.content.Context
+import android.os.Process.myUserHandle
+import android.os.UserHandle
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.WidgetPredictionsRequester.buildBundleForPredictionSession
+import com.android.launcher3.model.WidgetPredictionsRequester.filterPredictions
+import com.android.launcher3.model.WidgetPredictionsRequester.notOnUiSurfaceFilter
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.google.common.truth.Truth.assertThat
+import java.util.function.Predicate
+import junit.framework.Assert.assertNotNull
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+
+@RunWith(AndroidJUnit4::class)
+class WidgetsPredictionsRequesterTest {
+
+    private lateinit var mUserHandle: UserHandle
+    private lateinit var context: Context
+    private lateinit var deviceProfile: DeviceProfile
+    private lateinit var testInvariantProfile: InvariantDeviceProfile
+
+    private lateinit var widget1aInfo: AppWidgetProviderInfo
+    private lateinit var widget1bInfo: AppWidgetProviderInfo
+    private lateinit var widget2Info: AppWidgetProviderInfo
+
+    private lateinit var widgetItem1a: WidgetItem
+    private lateinit var widgetItem1b: WidgetItem
+    private lateinit var widgetItem2: WidgetItem
+
+    private lateinit var allWidgets: Map<PackageUserKey, List<WidgetItem>>
+
+    @Mock private lateinit var iconCache: IconCache
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        mUserHandle = myUserHandle()
+        context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+        testInvariantProfile = LauncherAppState.getIDP(context)
+        deviceProfile = testInvariantProfile.getDeviceProfile(context).copy(context)
+
+        widget1aInfo =
+            createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_1_PACKAGE_NAME, APP_1_PROVIDER_A_CLASS_NAME)
+            )
+        widget1bInfo =
+            createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_1_PACKAGE_NAME, APP_1_PROVIDER_B_CLASS_NAME)
+            )
+        widgetItem1a = createWidgetItem(widget1aInfo)
+        widgetItem1b = createWidgetItem(widget1bInfo)
+
+        widget2Info =
+            createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_2_PACKAGE_NAME, APP_2_PROVIDER_1_CLASS_NAME)
+            )
+        widgetItem2 = createWidgetItem(widget2Info)
+
+        allWidgets =
+            mapOf(
+                PackageUserKey(APP_1_PACKAGE_NAME, mUserHandle) to
+                    listOf(widgetItem1a, widgetItem1b),
+                PackageUserKey(APP_2_PACKAGE_NAME, mUserHandle) to listOf(widgetItem2),
+            )
+    }
+
+    @Test
+    fun buildBundleForPredictionSession_includesAddedAppWidgets() {
+        val existingWidgets = arrayListOf(widget1aInfo, widget1bInfo, widget2Info)
+
+        val bundle = buildBundleForPredictionSession(existingWidgets, TEST_UI_SURFACE)
+        val addedWidgetsBundleExtra =
+            bundle.getParcelableArrayList(BUNDLE_KEY_ADDED_APP_WIDGETS, AppTarget::class.java)
+
+        assertNotNull(addedWidgetsBundleExtra)
+        assertThat(addedWidgetsBundleExtra)
+            .containsExactly(
+                buildExpectedAppTargetEvent(
+                    /*pkg=*/ APP_1_PACKAGE_NAME,
+                    /*providerClassName=*/ APP_1_PROVIDER_A_CLASS_NAME,
+                    /*user=*/ mUserHandle
+                ),
+                buildExpectedAppTargetEvent(
+                    /*pkg=*/ APP_1_PACKAGE_NAME,
+                    /*providerClassName=*/ APP_1_PROVIDER_B_CLASS_NAME,
+                    /*user=*/ mUserHandle
+                ),
+                buildExpectedAppTargetEvent(
+                    /*pkg=*/ APP_2_PACKAGE_NAME,
+                    /*providerClassName=*/ APP_2_PROVIDER_1_CLASS_NAME,
+                    /*user=*/ mUserHandle
+                )
+            )
+    }
+
+    @Test
+    fun filterPredictions_notOnUiSurfaceFilter_returnsOnlyEligiblePredictions() {
+        val widgetsAlreadyOnSurface = arrayListOf(widget1bInfo)
+        val filter: Predicate<WidgetItem> = notOnUiSurfaceFilter(widgetsAlreadyOnSurface)
+
+        val predictions =
+            listOf(
+                // already on surface
+                AppTarget(
+                    AppTargetId(APP_1_PACKAGE_NAME),
+                    APP_1_PACKAGE_NAME,
+                    APP_1_PROVIDER_B_CLASS_NAME,
+                    mUserHandle
+                ),
+                // eligible
+                AppTarget(
+                    AppTargetId(APP_2_PACKAGE_NAME),
+                    APP_2_PACKAGE_NAME,
+                    APP_2_PROVIDER_1_CLASS_NAME,
+                    mUserHandle
+                )
+            )
+
+        // only 2 was eligible
+        assertThat(filterPredictions(predictions, allWidgets, filter)).containsExactly(widgetItem2)
+    }
+
+    @Test
+    fun filterPredictions_appPredictions_returnsWidgetFromPackage() {
+        val widgetsAlreadyOnSurface = arrayListOf(widget1bInfo)
+        val filter: Predicate<WidgetItem> = notOnUiSurfaceFilter(widgetsAlreadyOnSurface)
+
+        val predictions =
+            listOf(
+                AppTarget(
+                    AppTargetId(APP_1_PACKAGE_NAME),
+                    APP_1_PACKAGE_NAME,
+                    "$APP_1_PACKAGE_NAME.SomeActivity",
+                    mUserHandle
+                ),
+                AppTarget(
+                    AppTargetId(APP_2_PACKAGE_NAME),
+                    APP_2_PACKAGE_NAME,
+                    "$APP_2_PACKAGE_NAME.SomeActivity2",
+                    mUserHandle
+                ),
+            )
+
+        assertThat(filterPredictions(predictions, allWidgets, filter))
+            .containsExactly(widgetItem1a, widgetItem2)
+    }
+
+    private fun createWidgetItem(
+        providerInfo: AppWidgetProviderInfo,
+    ): WidgetItem {
+        val widgetInfo = LauncherAppWidgetProviderInfo.fromProviderInfo(context, providerInfo)
+        return WidgetItem(widgetInfo, testInvariantProfile, iconCache, context)
+    }
+
+    companion object {
+        const val TEST_UI_SURFACE = "widgets_test"
+        const val BUNDLE_KEY_ADDED_APP_WIDGETS = "added_app_widgets"
+
+        const val APP_1_PACKAGE_NAME = "com.example.app1"
+        const val APP_1_PROVIDER_A_CLASS_NAME = "app1Provider1"
+        const val APP_1_PROVIDER_B_CLASS_NAME = "app1Provider2"
+
+        const val APP_2_PACKAGE_NAME = "com.example.app2"
+        const val APP_2_PROVIDER_1_CLASS_NAME = "app2Provider1"
+
+        const val TEST_PACKAGE = "pkg"
+
+        private fun buildExpectedAppTargetEvent(
+            pkg: String,
+            providerClassName: String,
+            userHandle: UserHandle
+        ): AppTargetEvent {
+            val appTarget =
+                AppTarget.Builder(
+                        /*id=*/ AppTargetId("widget:$pkg"),
+                        /*packageName=*/ pkg,
+                        /*user=*/ userHandle
+                    )
+                    .setClassName(providerClassName)
+                    .build()
+            return AppTargetEvent.Builder(appTarget, AppTargetEvent.ACTION_PIN)
+                .setLaunchLocation(TEST_UI_SURFACE)
+                .build()
+        }
+    }
+}