diff --git a/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java b/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
index 34316db..023486f 100644
--- a/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
+++ b/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.appprediction;
 
+import static com.android.launcher3.BubbleTextView.DISPLAY_PREDICTION_ROW;
+
 import android.annotation.TargetApi;
 import android.content.Context;
 import android.graphics.Canvas;
@@ -36,7 +38,6 @@
 import com.android.launcher3.allapps.FloatingHeaderRow;
 import com.android.launcher3.allapps.FloatingHeaderView;
 import com.android.launcher3.anim.AlphaUpdateListener;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.keyboard.FocusIndicatorHelper;
 import com.android.launcher3.keyboard.FocusIndicatorHelper.SimpleFocusIndicatorHelper;
 import com.android.launcher3.model.data.ItemInfo;
@@ -106,22 +107,12 @@
                 mActivityContext.getAppsView().getAppsStore().unregisterIconContainer(this);
             }
         }
-
-        // Set the predicted row in All Apps' text line to 1.
-        if (FeatureFlags.ENABLE_TWOLINE_ALLAPPS.get()
-                || FeatureFlags.ENABLE_TWOLINE_DEVICESEARCH.get()) {
-            for (int i = 0; i < getChildCount(); i++) {
-                BubbleTextView icon = (BubbleTextView) getChildAt(i);
-                icon.setMaxLines(1);
-            }
-        }
     }
 
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
         super.onMeasure(widthMeasureSpec, MeasureSpec.makeMeasureSpec(getExpectedHeight(),
                 MeasureSpec.EXACTLY));
-        updateVisibility();
     }
 
     @Override
@@ -231,6 +222,7 @@
             icon.reset();
             if (predictionCount > i) {
                 icon.setVisibility(View.VISIBLE);
+                icon.setDisplay(DISPLAY_PREDICTION_ROW);
                 icon.applyFromWorkspaceItem(mPredictedApps.get(i));
             } else {
                 icon.setVisibility(predictionCount == 0 ? GONE : INVISIBLE);
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 76a1239..73e392d 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -85,6 +85,7 @@
             <enum name="taskbar" value="5" />
             <enum name="search_result_tall" value="6" />
             <enum name="search_result_small" value="7" />
+            <enum name="prediction_row" value="8" />
         </attr>
         <attr name="centerVertically" format="boolean" />
     </declare-styleable>
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 05d434e..2356bcc 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -94,11 +94,12 @@
         IconLabelDotView, DraggableView, Reorderable {
 
     private static final int DISPLAY_WORKSPACE = 0;
-    private static final int DISPLAY_ALL_APPS = 1;
+    public static final int DISPLAY_ALL_APPS = 1;
     private static final int DISPLAY_FOLDER = 2;
     protected static final int DISPLAY_TASKBAR = 5;
     private static final int DISPLAY_SEARCH_RESULT = 6;
     private static final int DISPLAY_SEARCH_RESULT_SMALL = 7;
+    public static final int DISPLAY_PREDICTION_ROW = 8;
 
     private static final float MIN_LETTER_SPACING = -0.05f;
     private static final int MAX_SEARCH_LOOP_COUNT = 20;
@@ -211,7 +212,7 @@
             setCompoundDrawablePadding(grid.iconDrawablePaddingPx);
             defaultIconSize = grid.iconSizePx;
             setCenterVertically(grid.iconCenterVertically);
-        } else if (mDisplay == DISPLAY_ALL_APPS) {
+        } else if (mDisplay == DISPLAY_ALL_APPS || mDisplay == DISPLAY_PREDICTION_ROW) {
             setTextSize(TypedValue.COMPLEX_UNIT_PX, grid.allAppsIconTextSizePx);
             setCompoundDrawablePadding(grid.allAppsIconDrawablePaddingPx);
             defaultIconSize = grid.allAppsIconSizePx;
@@ -402,7 +403,7 @@
      *  Only if actual text can be displayed in two line, the {@code true} value will be effective.
      */
     protected boolean shouldUseTwoLine() {
-        return (FeatureFlags.ENABLE_TWOLINE_ALLAPPS.get() && mDisplay == DISPLAY_ALL_APPS)
+        return  (FeatureFlags.ENABLE_TWOLINE_ALLAPPS.get() && mDisplay == DISPLAY_ALL_APPS)
                 || (FeatureFlags.ENABLE_TWOLINE_DEVICESEARCH.get()
                 && mDisplay == DISPLAY_SEARCH_RESULT);
     }
@@ -424,10 +425,10 @@
         }
     }
 
-    /** This is used for testing to forcefully set the display to ALL_APPS */
+    /** This is used for testing to forcefully set the display. */
     @VisibleForTesting
-    public void setDisplayAllApps() {
-        mDisplay = DISPLAY_ALL_APPS;
+    public void setDisplay(int display) {
+        mDisplay = display;
     }
 
     /**
diff --git a/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java b/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
index fdba4eb..ba17fdc 100644
--- a/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
+++ b/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
@@ -18,6 +18,8 @@
 
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 
+import static com.android.launcher3.BubbleTextView.DISPLAY_ALL_APPS;
+import static com.android.launcher3.BubbleTextView.DISPLAY_PREDICTION_ROW;
 import static com.android.launcher3.config.FeatureFlags.ENABLE_TWOLINE_ALLAPPS;
 
 import static org.junit.Assert.assertEquals;
@@ -79,7 +81,6 @@
         mContext = new ActivityContextWrapper(getApplicationContext());
         mBubbleTextView = new BubbleTextView(mContext);
         mBubbleTextView.reset();
-        mBubbleTextView.setDisplayAllApps();
 
         BubbleTextView testView = new BubbleTextView(mContext);
         testView.setTypeface(Typeface.MONOSPACE);
@@ -104,6 +105,7 @@
     public void testEmptyString_flagOn() {
         try (AutoCloseable flag = TestUtil.overrideFlag(ENABLE_TWOLINE_ALLAPPS, true)) {
             mItemInfoWithIcon.title = EMPTY_STRING;
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
@@ -118,6 +120,7 @@
     public void testEmptyString_flagOff() {
         try (AutoCloseable flag = TestUtil.overrideFlag(ENABLE_TWOLINE_ALLAPPS, false)) {
             mItemInfoWithIcon.title = EMPTY_STRING;
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
@@ -134,6 +137,7 @@
             // test string: "Battery Stats"
             mItemInfoWithIcon.title = TEST_STRING_WITH_SPACE_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -149,6 +153,7 @@
             // test string: "Battery Stats"
             mItemInfoWithIcon.title = TEST_STRING_WITH_SPACE_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -164,6 +169,7 @@
             // test string: "flutterappflorafy"
             mItemInfoWithIcon.title = TEST_LONG_STRING_NO_SPACE_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -179,6 +185,7 @@
             // test string: "flutterappflorafy"
             mItemInfoWithIcon.title = TEST_LONG_STRING_NO_SPACE_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -194,6 +201,7 @@
             // test string: "System UWB Field Test"
             mItemInfoWithIcon.title = TEST_LONG_STRING_WITH_SPACE_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -209,6 +217,7 @@
             // test string: "System UWB Field Test"
             mItemInfoWithIcon.title = TEST_LONG_STRING_WITH_SPACE_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -224,6 +233,7 @@
             // test string: "LEGO®Builder"
             mItemInfoWithIcon.title = TEST_LONG_STRING_SYMBOL_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -239,6 +249,7 @@
             // test string: "LEGO®Builder"
             mItemInfoWithIcon.title = TEST_LONG_STRING_SYMBOL_LONGER_THAN_CHAR_LIMIT;
             mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setDisplay(DISPLAY_ALL_APPS);
             mBubbleTextView.setTypeface(Typeface.MONOSPACE);
             mBubbleTextView.measure(mLimitedWidth, 0);
             mBubbleTextView.onPreDraw();
@@ -291,4 +302,20 @@
                 breakPoints);
         assertEquals(TEST_LONG_STRING_SYMBOL_LONGER_THAN_CHAR_LIMIT_RESULT, newString);
     }
+
+    @Test
+    public void testEnsurePredictionRowIsOneLine() {
+        try (AutoCloseable flag = TestUtil.overrideFlag(ENABLE_TWOLINE_ALLAPPS, true)) {
+            // test string: "Battery Stats"
+            mItemInfoWithIcon.title = TEST_STRING_WITH_SPACE_LONGER_THAN_CHAR_LIMIT;
+            mBubbleTextView.setDisplay(DISPLAY_PREDICTION_ROW);
+            mBubbleTextView.applyLabel(mItemInfoWithIcon);
+            mBubbleTextView.setTypeface(Typeface.MONOSPACE);
+            mBubbleTextView.measure(mLimitedWidth, 0);
+            mBubbleTextView.onPreDraw();
+            assertEquals(ONE_LINE, mBubbleTextView.getLineCount());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
 }