diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index 3769124..eb3280e2 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -398,3 +398,10 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "enable_dismiss_prediction_undo"
+    namespace: "launcher"
+    description: "Show an 'Undo' snackbar when users dismiss a predicted hotseat item"
+    bug: "270394476"
+}
diff --git a/src/com/android/launcher3/DropTargetHandler.kt b/src/com/android/launcher3/DropTargetHandler.kt
index e022159..f1029b1 100644
--- a/src/com/android/launcher3/DropTargetHandler.kt
+++ b/src/com/android/launcher3/DropTargetHandler.kt
@@ -35,8 +35,7 @@
                     target?.let {
                         deferred.mPackageName = it.packageName
                         mLauncher.addEventCallback(EVENT_RESUMED) { deferred.onLauncherResume() }
-                    }
-                        ?: deferred.sendFailure()
+                    } ?: deferred.sendFailure()
                 }
             }
         }
@@ -47,19 +46,10 @@
         mLauncher.appWidgetHolder.startConfigActivity(
             mLauncher,
             widgetId,
-            ActivityCodes.REQUEST_RECONFIGURE_APPWIDGET
+            ActivityCodes.REQUEST_RECONFIGURE_APPWIDGET,
         )
     }
 
-    fun dismissPrediction(
-        announcement: CharSequence,
-        onActionClicked: Runnable,
-        onDismiss: Runnable?
-    ) {
-        mLauncher.dragLayer.announceForAccessibility(announcement)
-        Snackbar.show(mLauncher, R.string.item_removed, R.string.undo, onDismiss, onActionClicked)
-    }
-
     fun getViewUnderDrag(info: ItemInfo): View? {
         return if (
             info is LauncherAppWidgetInfo &&
@@ -95,7 +85,7 @@
             R.string.item_removed,
             R.string.undo,
             mLauncher.modelWriter::commitDelete,
-            onUndoClicked
+            onUndoClicked,
         )
     }
 
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index 0a4fb73..8d1e61f 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -7,7 +7,6 @@
 import static com.android.launcher3.accessibility.LauncherAccessibilityDelegate.INVALID;
 import static com.android.launcher3.accessibility.LauncherAccessibilityDelegate.RECONFIGURE;
 import static com.android.launcher3.accessibility.LauncherAccessibilityDelegate.UNINSTALL;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_DISMISS_PREDICTION_UNDO;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ITEM_DROPPED_ON_DONT_SUGGEST;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ITEM_DROPPED_ON_UNINSTALL;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ITEM_UNINSTALL_CANCELLED;
@@ -36,7 +35,6 @@
 
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.logging.InstanceId;
@@ -242,8 +240,7 @@
 
     @Override
     public void completeDrop(final DragObject d) {
-        ComponentName target = performDropAction(getViewUnderDrag(d.dragInfo), d.dragInfo,
-                d.logInstanceId);
+        ComponentName target = performDropAction(getViewUnderDrag(d.dragInfo), d.dragInfo);
         mDropTargetHandler.onSecondaryTargetCompleteDrop(target, d);
     }
 
@@ -275,7 +272,7 @@
      * Performs the drop action and returns the target component for the dragObject or null if
      * the action was not performed.
      */
-    protected ComponentName performDropAction(View view, ItemInfo info, InstanceId instanceId) {
+    protected ComponentName performDropAction(View view, ItemInfo info) {
         if (mCurrentAccessibilityAction == RECONFIGURE) {
             int widgetId = getReconfigurableWidgetId(view);
             if (widgetId != INVALID_APPWIDGET_ID) {
@@ -283,21 +280,6 @@
             }
             return null;
         }
-        if (mCurrentAccessibilityAction == DISMISS_PREDICTION) {
-            if (FeatureFlags.ENABLE_DISMISS_PREDICTION_UNDO.get()) {
-                CharSequence announcement = getContext().getString(R.string.item_removed);
-                mDropTargetHandler
-                        .dismissPrediction(announcement, () -> {
-                        }, () -> {
-                            mStatsLogManager.logger()
-                                    .withInstanceId(instanceId)
-                                    .withItemInfo(info)
-                                    .log(LAUNCHER_DISMISS_PREDICTION_UNDO);
-                        });
-            }
-            return null;
-        }
-
         return performUninstall(getContext(), getUninstallTarget(getContext(), info), info);
     }
 
@@ -332,9 +314,8 @@
 
     @Override
     public void onAccessibilityDrop(View view, ItemInfo item) {
-        InstanceId instanceId = new InstanceIdSequence().newInstanceId();
-        doLog(instanceId, item);
-        performDropAction(view, item, instanceId);
+        doLog(new InstanceIdSequence().newInstanceId(), item);
+        performDropAction(view, item);
     }
 
     /**
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index eb65320..998b2bb 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -63,10 +63,6 @@
      * <p>
      */
     // TODO(Block 3): Clean up flags
-    public static final BooleanFlag ENABLE_DISMISS_PREDICTION_UNDO = getDebugFlag(270394476,
-            "ENABLE_DISMISS_PREDICTION_UNDO", DISABLED,
-            "Show an 'Undo' snackbar when users dismiss a predicted hotseat item");
-
     public static final BooleanFlag ENABLE_WORKSPACE_LOADING_OPTIMIZATION = getDebugFlag(251502424,
             "ENABLE_WORKSPACE_LOADING_OPTIMIZATION", DISABLED,
             "load the current workspace screen visible to the user before the rest rather than "
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index 1b245ab..63c9d94 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -1,5 +1,6 @@
 package com.android.launcher3.popup;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_DISMISS_PREDICTION_UNDO;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_INSTALL_SYSTEM_SHORTCUT_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_UNINSTALL_SYSTEM_SHORTCUT_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_APP_INFO_TAP;
@@ -41,6 +42,7 @@
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.Snackbar;
 import com.android.launcher3.widget.WidgetsBottomSheet;
 import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 
@@ -336,6 +338,14 @@
             mTarget.getStatsLogManager().logger()
                     .withItemInfo(mItemInfo)
                     .log(LAUNCHER_SYSTEM_SHORTCUT_DONT_SUGGEST_APP_TAP);
+            if (Flags.enableDismissPredictionUndo()) {
+                Snackbar.show(mTarget,
+                        view.getContext().getString(R.string.item_removed), R.string.undo,
+                        () -> { }, () ->
+                            mTarget.getStatsLogManager().logger()
+                                    .withItemInfo(mItemInfo)
+                                    .log(LAUNCHER_DISMISS_PREDICTION_UNDO));
+            }
         }
     }
 
diff --git a/tests/src/com/android/launcher3/popup/SystemShortcutTest.java b/tests/src/com/android/launcher3/popup/SystemShortcutTest.java
index dcfcad5..f54668c 100644
--- a/tests/src/com/android/launcher3/popup/SystemShortcutTest.java
+++ b/tests/src/com/android/launcher3/popup/SystemShortcutTest.java
@@ -19,19 +19,26 @@
 import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.launcher3.AbstractFloatingView.TYPE_SNACKBAR;
+import static com.android.launcher3.Flags.FLAG_ENABLE_DISMISS_PREDICTION_UNDO;
 import static com.android.launcher3.Flags.FLAG_ENABLE_PRIVATE_SPACE;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_ALL_APPS;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_DISMISS_PREDICTION_UNDO;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_DONT_SUGGEST_APP_TAP;
 import static com.android.launcher3.model.data.WorkspaceItemInfo.FLAG_SUPPORTS_WEB_UI;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -50,8 +57,13 @@
 
 import androidx.test.annotation.UiThreadTest;
 import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.launcher3.AbstractFloatingView;
+import com.android.launcher3.R;
 import com.android.launcher3.allapps.PrivateProfileManager;
+import com.android.launcher3.logging.StatsLogManager;
+import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -61,8 +73,9 @@
 import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext;
 import com.android.launcher3.util.LauncherMultivalentJUnit;
 import com.android.launcher3.util.TestSandboxModelContextWrapper;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.UserIconInfo;
-import com.android.launcher3.views.BaseDragLayer;
+import com.android.launcher3.views.Snackbar;
 import com.android.launcher3.widget.picker.model.WidgetPickerDataProvider;
 import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 
@@ -72,6 +85,7 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -88,25 +102,31 @@
     private PrivateProfileManager mPrivateProfileManager;
     private WidgetPickerDataProvider mWidgetPickerDataProvider;
     private AppInfo mAppInfo;
+
     @Mock UserCache mUserCache;
     @Mock ApiWrapper mApiWrapper;
-    @Mock BaseDragLayer mBaseDragLayer;
     @Mock UserIconInfo mUserIconInfo;
     @Mock LauncherActivityInfo mLauncherActivityInfo;
     @Mock ApplicationInfo mApplicationInfo;
     @Mock Intent mIntent;
+    @Mock StatsLogManager mStatsLogManager;
+    @Mock(answer = Answers.RETURNS_SELF) StatsLogger mStatsLogger;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mSandboxContext.putObject(UserCache.INSTANCE, mUserCache);
         mSandboxContext.putObject(ApiWrapper.INSTANCE, mApiWrapper);
-        mTestContext = new TestSandboxModelContextWrapper(mSandboxContext);
-        mView = new View(mSandboxContext);
-        spyOn(mTestContext);
+        mTestContext = new TestSandboxModelContextWrapper(mSandboxContext) {
+            @Override
+            public StatsLogManager getStatsLogManager() {
+                return mStatsLogManager;
+            }
+        };
         spyOn(mSandboxContext);
-        doReturn(mBaseDragLayer).when(mTestContext).getDragLayer();
+        doReturn(mStatsLogger).when(mStatsLogManager).logger();
 
+        mView = new View(mTestContext);
         mItemInfo = new ItemInfo();
 
         LauncherApps mLauncherApps = mSandboxContext.spyService(LauncherApps.class);
@@ -114,7 +134,6 @@
         when(mLauncherActivityInfo.getApplicationInfo()).thenReturn(mApplicationInfo);
 
         when(mUserCache.getUserInfo(any())).thenReturn(mUserIconInfo);
-        when(mBaseDragLayer.getChildCount()).thenReturn(0);
         mPrivateProfileManager = mTestContext.getAppsView().getPrivateProfileManager();
         spyOn(mPrivateProfileManager);
         when(mPrivateProfileManager.getProfileUser()).thenReturn(PRIVATE_HANDLE);
@@ -168,6 +187,7 @@
     }
 
     @Test
+    @DisableFlags(FLAG_ENABLE_DISMISS_PREDICTION_UNDO)
     public void testDontSuggestAppForPredictedItem() {
         mAppInfo = new AppInfo();
         mAppInfo.componentName = new ComponentName(mTestContext, getClass());
@@ -176,7 +196,36 @@
         SystemShortcut systemShortcut = SystemShortcut.DONT_SUGGEST_APP
                 .getShortcut(mTestContext, mAppInfo, mView);
         assertNotNull(systemShortcut);
-        systemShortcut.onClick(mView);
+
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, () -> systemShortcut.onClick(mView));
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+
+        verify(mStatsLogger).log(eq(LAUNCHER_SYSTEM_SHORTCUT_DONT_SUGGEST_APP_TAP));
+        assertFalse(AbstractFloatingView.hasOpenView(mTestContext, TYPE_SNACKBAR));
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DISMISS_PREDICTION_UNDO)
+    public void testDontSuggestAppForPredictedItemWithUndo() {
+        mAppInfo = new AppInfo();
+        mAppInfo.componentName = new ComponentName(mTestContext, getClass());
+        mAppInfo.container = CONTAINER_HOTSEAT_PREDICTION;
+        assertTrue(mAppInfo.isPredictedItem());
+        SystemShortcut systemShortcut = SystemShortcut.DONT_SUGGEST_APP
+                .getShortcut(mTestContext, mAppInfo, mView);
+        assertNotNull(systemShortcut);
+
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, () -> systemShortcut.onClick(mView));
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+        verify(mStatsLogger).log(eq(LAUNCHER_SYSTEM_SHORTCUT_DONT_SUGGEST_APP_TAP));
+
+        // Undo bar shown
+        Snackbar snackbar = AbstractFloatingView.getOpenView(mTestContext, TYPE_SNACKBAR);
+        assertNotNull(snackbar);
+        reset(mStatsLogger);
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, snackbar.findViewById(
+                R.id.action)::performClick);
+        verify(mStatsLogger).log(eq(LAUNCHER_DISMISS_PREDICTION_UNDO));
     }
 
     @Test