diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 39f2c00c99..64bb05e1ab 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -19,6 +19,7 @@ import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 
 import android.app.prediction.AppTarget;
+import android.content.ComponentName;
 import android.content.Context;
 import android.text.TextUtils;
 
@@ -29,7 +30,6 @@ import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
 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;
 
@@ -64,36 +64,24 @@ public final class WidgetsPredictionUpdateTask implements ModelUpdateTask {
                 widget -> new ComponentKey(widget.providerName, widget.user)).collect(
                 Collectors.toSet());
         Predicate<WidgetItem> notOnWorkspace = w -> !widgetsInWorkspace.contains(w);
-        Map<PackageUserKey, List<WidgetItem>> allWidgets =
-                dataModel.widgetsModel.getAllWidgetsWithoutShortcuts();
+        Map<ComponentKey, WidgetItem> allWidgets =
+                dataModel.widgetsModel.getAllWidgetComponentsWithoutShortcuts();
 
         List<WidgetItem> servicePredictedItems = new ArrayList<>();
-        List<WidgetItem> localFilteredWidgets = new ArrayList<>();
 
         for (AppTarget app : mTargets) {
-            PackageUserKey packageUserKey = new PackageUserKey(app.getPackageName(), app.getUser());
-            List<WidgetItem> widgets = allWidgets.get(packageUserKey);
-            if (widgets == null || widgets.isEmpty()) {
+            ComponentKey componentKey = new ComponentKey(
+                    new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
+            WidgetItem widget = allWidgets.get(componentKey);
+            if (widget == null) {
                 continue;
             }
             String className = app.getClassName();
             if (!TextUtils.isEmpty(className)) {
-                WidgetItem item = widgets.stream()
-                        .filter(w -> className.equals(w.componentName.getClassName()))
-                        .filter(notOnWorkspace)
-                        .findFirst()
-                        .orElse(null);
-                if (item != null) {
-                    servicePredictedItems.add(item);
-                    continue;
+                if (notOnWorkspace.test(widget)) {
+                    servicePredictedItems.add(widget);
                 }
             }
-            // No widget was added by the service, try local filtering
-            widgets.stream().filter(notOnWorkspace).findFirst()
-                    .ifPresent(localFilteredWidgets::add);
-        }
-        if (servicePredictedItems.isEmpty()) {
-            servicePredictedItems.addAll(localFilteredWidgets);
         }
 
         List<ItemInfo> items;
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 8702f70dfe..7b57c81b74 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -184,7 +184,7 @@ public final class WidgetsPredicationUpdateTaskTest {
     }
 
     @Test
-    public void widgetsRecommendationRan_shouldReturnPackageWidgetsWhenEmpty() {
+    public void widgetsRecommendationRan_shouldReturnEmptyWidgetsWhenEmpty() {
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
 
             // Not installed widget
@@ -204,19 +204,12 @@ public final class WidgetsPredicationUpdateTaskTest {
                     newWidgetsPredicationTask(List.of(widget5, widget3, widget4, widget1)));
             runOnExecutorSync(MAIN_EXECUTOR, () -> { });
 
-            // THEN only 2 widgets are returned because the launcher only filters out
-            // non-exist widgets.
+            // Only widgets suggested by prediction system are returned.
             List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
                     .stream()
                     .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
                     .collect(Collectors.toList());
-            assertThat(recommendedWidgets).hasSize(2);
-            recommendedWidgets.forEach(pendingAddWidgetInfo ->
-                    assertThat(pendingAddWidgetInfo.recommendationCategory).isNotNull()
-            );
-            // Another widget from the same package
-            assertWidgetInfo(recommendedWidgets.get(0).info, mApp4Provider2);
-            assertWidgetInfo(recommendedWidgets.get(1).info, mApp1Provider1);
+            assertThat(recommendedWidgets).hasSize(0);
         });
     }
 
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 5e0edb3139..4cba0b5239 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -32,6 +32,7 @@ import com.android.launcher3.icons.ComponentWithLabelAndIcon;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
+import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
@@ -129,6 +130,22 @@ public class WidgetsModel {
         return packagesToWidgets;
     }
 
+    /**
+     * Returns a map of widget component keys to corresponding widget items. Excludes the
+     * shortcuts.
+     */
+    public synchronized Map<ComponentKey, WidgetItem> getAllWidgetComponentsWithoutShortcuts() {
+        if (!WIDGETS_ENABLED) {
+            return Collections.emptyMap();
+        }
+        Map<ComponentKey, WidgetItem> widgetsMap = new HashMap<>();
+        mWidgetsList.forEach((packageItemInfo, widgetsAndShortcuts) ->
+                widgetsAndShortcuts.stream().filter(item -> item.widgetInfo != null).forEach(
+                        item -> widgetsMap.put(new ComponentKey(item.componentName, item.user),
+                                item)));
+        return widgetsMap;
+    }
+
     /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.