diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java
index bd78573..0b8c1c5 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionAppTracker.java
@@ -15,7 +15,7 @@
  */
 package com.android.launcher3.appprediction;
 
-import static com.android.launcher3.appprediction.PredictionUiStateManager.KEY_APP_SUGGESTION;
+import static com.android.launcher3.InvariantDeviceProfile.CHANGE_FLAG_GRID;
 
 import android.annotation.TargetApi;
 import android.app.prediction.AppPredictionContext;
@@ -26,8 +26,6 @@
 import android.app.prediction.AppTargetId;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.SharedPreferences;
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener;
 import android.os.Build;
 import android.os.Handler;
 import android.os.Message;
@@ -35,12 +33,10 @@
 import android.util.Log;
 
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.Utilities;
+import com.android.launcher3.appprediction.PredictionUiStateManager.Client;
 import com.android.launcher3.model.AppLaunchTracker;
 import com.android.launcher3.util.UiThreadHelper;
 
-import com.android.launcher3.appprediction.PredictionUiStateManager.Client;
-
 import androidx.annotation.UiThread;
 import androidx.annotation.WorkerThread;
 
@@ -48,8 +44,7 @@
  * Subclass of app tracker which publishes the data to the prediction engine and gets back results.
  */
 @TargetApi(Build.VERSION_CODES.Q)
-public class PredictionAppTracker extends AppLaunchTracker
-        implements OnSharedPreferenceChangeListener {
+public class PredictionAppTracker extends AppLaunchTracker {
 
     private static final String TAG = "PredictionAppTracker";
     private static final boolean DBG = false;
@@ -62,8 +57,6 @@
     private final Context mContext;
     private final Handler mMessageHandler;
 
-    private boolean mEnabled;
-
     // Accessed only on worker thread
     private AppPredictor mHomeAppPredictor;
     private AppPredictor mRecentsOverviewPredictor;
@@ -71,24 +64,16 @@
     public PredictionAppTracker(Context context) {
         mContext = context;
         mMessageHandler = new Handler(UiThreadHelper.getBackgroundLooper(), this::handleMessage);
-
-        SharedPreferences prefs = Utilities.getPrefs(context);
-        setEnabled(prefs.getBoolean(KEY_APP_SUGGESTION, true));
-        prefs.registerOnSharedPreferenceChangeListener(this);
         InvariantDeviceProfile.INSTANCE.get(mContext).addOnChangeListener(this::onIdpChanged);
+
+        mMessageHandler.sendEmptyMessage(MSG_INIT);
     }
 
     @UiThread
     private void onIdpChanged(int changeFlags, InvariantDeviceProfile profile) {
-        // Reinitialize everything
-        setEnabled(mEnabled);
-    }
-
-    @Override
-    @UiThread
-    public void onSharedPreferenceChanged(SharedPreferences prefs, String key) {
-        if (KEY_APP_SUGGESTION.equals(key)) {
-            setEnabled(prefs.getBoolean(KEY_APP_SUGGESTION, true));
+        if ((changeFlags & CHANGE_FLAG_GRID) != 0) {
+            // Reinitialize everything
+            mMessageHandler.sendEmptyMessage(MSG_INIT);
         }
     }
 
@@ -137,13 +122,13 @@
                 return true;
             }
             case MSG_LAUNCH: {
-                if (mEnabled && mHomeAppPredictor != null) {
+                if (mHomeAppPredictor != null) {
                     mHomeAppPredictor.notifyAppTargetEvent((AppTargetEvent) msg.obj);
                 }
                 return true;
             }
             case MSG_PREDICT: {
-                if (mEnabled && mHomeAppPredictor != null) {
+                if (mHomeAppPredictor != null) {
                     String client = (String) msg.obj;
                     if (Client.HOME.id.equals(client)) {
                         mHomeAppPredictor.requestPredictionUpdate();
@@ -168,18 +153,6 @@
         }
     }
 
-    @UiThread
-    public void setEnabled(boolean isEnabled) {
-        mEnabled = isEnabled;
-        if (isEnabled) {
-            mMessageHandler.removeMessages(MSG_DESTROY);
-            mMessageHandler.sendEmptyMessage(MSG_INIT);
-        } else {
-            mMessageHandler.removeMessages(MSG_INIT);
-            mMessageHandler.sendEmptyMessage(MSG_DESTROY);
-        }
-    }
-
     @Override
     @UiThread
     public void onStartShortcut(String packageName, String shortcutId, UserHandle user,
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java
index 54fd845..48a163d 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java
@@ -24,8 +24,7 @@
 import android.app.prediction.AppTarget;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.SharedPreferences;
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener;
+import android.os.Handler;
 import android.view.ViewTreeObserver.OnGlobalLayoutListener;
 
 import com.android.launcher3.AppInfo;
@@ -61,9 +60,10 @@
  * that client id.
  */
 public class PredictionUiStateManager implements OnGlobalLayoutListener, ItemInfoUpdateReceiver,
-        OnSharedPreferenceChangeListener, OnIDPChangeListener, OnUpdateListener {
+        OnIDPChangeListener, OnUpdateListener {
 
-    public static final String KEY_APP_SUGGESTION = "pref_show_predictions";
+    public static final String LAST_PREDICTION_ENABLED_STATE = "last_prediction_enabled_state";
+    private static final long INITIAL_CALLBACK_WAIT_TIMEOUT_MS = 5000;
 
     // TODO (b/129421797): Update the client constants
     public enum Client {
@@ -81,7 +81,6 @@
             new MainThreadInitializedObject<>(PredictionUiStateManager::new);
 
     private final Context mContext;
-    private final SharedPreferences mMainPrefs;
 
     private final DynamicItemCache mDynamicItemCache;
     private final List[] mPredictionServicePredictions;
@@ -94,9 +93,10 @@
     private PredictionState mPendingState;
     private PredictionState mCurrentState;
 
+    private boolean mGettingValidPredictionResults;
+
     private PredictionUiStateManager(Context context) {
         mContext = context;
-        mMainPrefs = Utilities.getPrefs(context);
 
         mDynamicItemCache = new DynamicItemCache(context, this::onAppsUpdated);
 
@@ -110,8 +110,14 @@
         for (int i = 0; i < mPredictionServicePredictions.length; i++) {
             mPredictionServicePredictions[i] = Collections.emptyList();
         }
-        // Listens for enable/disable signal, and predictions if using AiAi is disabled.
-        mMainPrefs.registerOnSharedPreferenceChangeListener(this);
+
+        mGettingValidPredictionResults = Utilities.getDevicePrefs(context)
+                .getBoolean(LAST_PREDICTION_ENABLED_STATE, true);
+        if (mGettingValidPredictionResults) {
+            new Handler().postDelayed(
+                    this::updatePredictionStateAfterCallback, INITIAL_CALLBACK_WAIT_TIMEOUT_MS);
+        }
+
         // Call this last
         mCurrentState = parseLastState();
     }
@@ -177,13 +183,6 @@
         }
     }
 
-    @Override
-    public void onSharedPreferenceChanged(SharedPreferences prefs, String key) {
-        if (KEY_APP_SUGGESTION.equals(key)) {
-            dispatchOnChange(true);
-        }
-    }
-
     private void applyState(PredictionState state) {
         boolean wasEnabled = mCurrentState.isEnabled;
         mCurrentState = state;
@@ -198,10 +197,24 @@
         }
     }
 
+    private void updatePredictionStateAfterCallback() {
+        boolean validResults = false;
+        for (List l : mPredictionServicePredictions) {
+            validResults |= l != null && !l.isEmpty();
+        }
+        if (validResults != mGettingValidPredictionResults) {
+            mGettingValidPredictionResults = validResults;
+            Utilities.getDevicePrefs(mContext).edit()
+                    .putBoolean(LAST_PREDICTION_ENABLED_STATE, true)
+                    .apply();
+        }
+        dispatchOnChange(true);
+    }
+
     public AppPredictor.Callback appPredictorCallback(Client client) {
         return targets -> {
             mPredictionServicePredictions[client.ordinal()] = targets;
-            dispatchOnChange(true);
+            updatePredictionStateAfterCallback();
         };
     }
 
@@ -217,7 +230,7 @@
 
     private PredictionState parseLastState() {
         PredictionState state = new PredictionState();
-        state.isEnabled = mMainPrefs.getBoolean(KEY_APP_SUGGESTION, true);
+        state.isEnabled = mGettingValidPredictionResults;
         if (!state.isEnabled) {
             state.apps = Collections.EMPTY_LIST;
             return state;
diff --git a/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java b/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java
index 72de80b..e6dbf87 100644
--- a/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java
+++ b/quickstep/tests/src/com/android/quickstep/AppPredictionsUITests.java
@@ -17,6 +17,7 @@
 package com.android.quickstep;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 
 import android.app.prediction.AppPredictor;
 import android.app.prediction.AppTarget;
@@ -25,7 +26,6 @@
 import android.content.pm.LauncherActivityInfo;
 import android.os.Process;
 import android.view.View;
-import android.widget.ProgressBar;
 
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.Launcher;
@@ -49,8 +49,7 @@
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class AppPredictionsUITests  extends AbstractQuickStepTest {
-    private static final String TAG = "AppPredictionsUITests";
+public class AppPredictionsUITests extends AbstractQuickStepTest {
 
     private LauncherActivityInfo mSampleApp1;
     private LauncherActivityInfo mSampleApp2;
@@ -91,20 +90,15 @@
         mActivityMonitor.startLauncher();
         mLauncher.pressHome().switchToAllApps();
 
-        // There has not been any update, verify that progress bar is showing
-        waitForLauncherCondition("Prediction is not in loading state", launcher -> {
-            ProgressBar p = findLoadingBar(launcher);
-            return p != null && p.isShown();
-        });
-
         // Dispatch an update
         sendPredictionUpdate(mSampleApp1, mSampleApp2);
+        // The first update should apply immediately.
         waitForLauncherCondition("Predictions were not updated in loading state",
                 launcher -> getPredictedApp(launcher).size() == 2);
     }
 
     /**
-     * Test tat prediction update is deferred if it is already visible
+     * Test that prediction update is deferred if it is already visible
      */
     @Test
     @Ignore // b/131188880
@@ -124,6 +118,19 @@
         assertEquals(3, getFromLauncher(this::getPredictedApp).size());
     }
 
+    @Test
+    public void testPredictionsDisabled() {
+        mActivityMonitor.startLauncher();
+        sendPredictionUpdate();
+        mLauncher.pressHome().switchToAllApps();
+
+        waitForLauncherCondition("Predictions were not updated in loading state",
+                launcher -> launcher.getAppsView().getFloatingHeaderView()
+                        .findFixedRowByType(PredictionRowView.class).getVisibility() == View.GONE);
+        assertFalse(PredictionUiStateManager.INSTANCE.get(mTargetContext)
+                .getCurrentState().isEnabled);
+    }
+
     public ArrayList<BubbleTextView> getPredictedApp(Launcher launcher) {
         PredictionRowView container = launcher.getAppsView().getFloatingHeaderView()
                 .findFixedRowByType(PredictionRowView.class);
@@ -138,20 +145,6 @@
         return predictedAppViews;
     }
 
-    private ProgressBar findLoadingBar(Launcher launcher) {
-        PredictionRowView container = launcher.getAppsView().getFloatingHeaderView()
-                .findFixedRowByType(PredictionRowView.class);
-
-        for (int i = 0; i < container.getChildCount(); i++) {
-            View view = container.getChildAt(i);
-            if (view instanceof ProgressBar) {
-                return (ProgressBar) view;
-            }
-        }
-        return null;
-    }
-
-
     private void sendPredictionUpdate(LauncherActivityInfo... activities) {
         getOnUiThread(() -> {
             List<AppTarget> targets = new ArrayList<>(activities.length);
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index b0ce5f5..4b6b3ee 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -64,7 +64,7 @@
     private void startTestApps() throws Exception {
         startAppFast(getAppPackageName());
         startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
-        startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CONTACTS));
+        startTestActivity(2);
 
         executeOnLauncher(launcher -> assertTrue(
                 "Launcher activity is the top activity; expecting another activity to be the top "
@@ -130,8 +130,8 @@
         OverviewTask task = mLauncher.pressHome().switchToOverview().getCurrentTask();
         assertNotNull("overview.getCurrentTask() returned null (1)", task);
         assertNotNull("OverviewTask.open returned null", task.open());
-        assertTrue("Contacts app didn't open from Overview", mDevice.wait(Until.hasObject(
-                By.pkg(resolveSystemApp(Intent.CATEGORY_APP_CONTACTS)).depth(0)),
+        assertTrue("Test activity didn't open from Overview", mDevice.wait(Until.hasObject(
+                By.pkg(getAppPackageName()).text("TestActivity2")),
                 LONG_WAIT_TIME_MS));
         executeOnLauncher(launcher -> assertTrue(
                 "Launcher activity is the top activity; expecting another activity to be the top "
diff --git a/tests/AndroidManifest-common.xml b/tests/AndroidManifest-common.xml
index 4051fbd..75ff66e 100644
--- a/tests/AndroidManifest-common.xml
+++ b/tests/AndroidManifest-common.xml
@@ -99,7 +99,8 @@
         </activity>
         <activity
             android:name="com.android.launcher3.testcomponent.BaseTestingActivity"
-            android:label="LauncherTestApp">
+            android:label="LauncherTestApp"
+            android:taskAffinity="com.android.launcher3.testcomponent.Affinity1">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN"/>
                 <category android:name="android.intent.category.LAUNCHER"/>
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index dedc6b3..43bdb9f 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -17,6 +17,7 @@
 
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.launcher3.ui.TaplTestsLauncher3.getAppPackageName;
 import static com.android.systemui.shared.system.QuickStepContract.NAV_BAR_MODE_2BUTTON_OVERLAY;
 import static com.android.systemui.shared.system.QuickStepContract.NAV_BAR_MODE_3BUTTON_OVERLAY;
 import static com.android.systemui.shared.system.QuickStepContract.NAV_BAR_MODE_GESTURAL_OVERLAY;
@@ -28,6 +29,7 @@
 
 import android.app.Instrumentation;
 import android.content.BroadcastReceiver;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
@@ -426,6 +428,22 @@
                 mDevice.wait(Until.hasObject(By.pkg(packageName).depth(0)), LONG_WAIT_TIME_MS));
     }
 
+    protected void startTestActivity(int activityNumber) {
+        final String packageName = getAppPackageName();
+        final Instrumentation instrumentation = getInstrumentation();
+        final Intent intent = instrumentation.getContext().getPackageManager().
+                getLaunchIntentForPackage(packageName);
+        intent.addCategory(Intent.CATEGORY_LAUNCHER);
+        intent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
+        intent.setComponent(new ComponentName(packageName,
+                "com.android.launcher3.tests.Activity" + activityNumber));
+        instrumentation.getTargetContext().startActivity(intent);
+        assertTrue(packageName + " didn't start",
+                mDevice.wait(
+                        Until.hasObject(By.pkg(packageName).text("TestActivity" + activityNumber)),
+                        LONG_WAIT_TIME_MS));
+    }
+
     protected static String resolveSystemApp(String category) {
         return getInstrumentation().getContext().getPackageManager().resolveActivity(
                 new Intent(Intent.ACTION_MAIN).addCategory(category),