diff --git a/Android.bp b/Android.bp
index 4e1c9d3..9d7aa73 100644
--- a/Android.bp
+++ b/Android.bp
@@ -64,11 +64,74 @@
 filegroup {
     name: "launcher-quickstep-src",
     srcs: [
-        "quickstep/src/**/*.java",
         "quickstep/src/**/*.kt",
+        "quickstep/src/**/*.java",
+        ":launcher-quickstep-processed-protolog-src",
     ],
 }
 
+// Launcher ProtoLog support
+filegroup {
+    name: "launcher-quickstep-unprocessed-protolog-src",
+    srcs: [
+        "quickstep/src_protolog/**/*.java",
+    ],
+}
+
+java_library {
+    name: "launcher-quickstep_protolog-groups",
+    srcs: [
+        "quickstep/src_protolog/**/*.java",
+    ],
+    static_libs: [
+        "protolog-group",
+        "androidx.annotation_annotation",
+        "com_android_launcher3_flags_lib",
+    ],
+}
+
+genrule {
+    name: "launcher-quickstep-processed-protolog-src",
+    srcs: [
+        ":protolog-impl",
+        ":launcher-quickstep-unprocessed-protolog-src",
+        ":launcher-quickstep_protolog-groups",
+    ],
+    tools: ["protologtool"],
+    cmd: "$(location protologtool) transform-protolog-calls " +
+        "--protolog-class com.android.internal.protolog.common.ProtoLog " +
+        "--loggroups-class com.android.quickstep.util.QuickstepProtoLogGroup " +
+        "--loggroups-jar $(location :launcher-quickstep_protolog-groups) " +
+        "--viewer-config-file-path /system_ext/etc/launcher.quickstep.protolog.pb " +
+        "--output-srcjar $(out) " +
+        "$(locations :launcher-quickstep-unprocessed-protolog-src)",
+    out: ["launcher.quickstep.protolog.srcjar"],
+}
+
+genrule {
+    name: "gen-launcher.quickstep.protolog.pb",
+    srcs: [
+        ":launcher-quickstep-unprocessed-protolog-src",
+        ":launcher-quickstep_protolog-groups",
+    ],
+    tools: ["protologtool"],
+    cmd: "$(location protologtool) generate-viewer-config " +
+        "--protolog-class com.android.internal.protolog.common.ProtoLog " +
+        "--loggroups-class com.android.quickstep.util.QuickstepProtoLogGroup " +
+        "--loggroups-jar $(location :launcher-quickstep_protolog-groups) " +
+        "--viewer-config-type proto " +
+        "--viewer-config $(out) " +
+        "$(locations :launcher-quickstep-unprocessed-protolog-src)",
+    out: ["launcher.quickstep.protolog.pb"],
+}
+
+prebuilt_etc {
+    name: "launcher.quickstep.protolog.pb",
+    system_ext_specific: true,
+    src: ":gen-launcher.quickstep.protolog.pb",
+    filename_from_src: true,
+}
+
 // Source code for quickstep dagger
 filegroup {
     name: "launcher-quickstep-dagger",
@@ -405,6 +468,7 @@
         "SystemUISharedLib",
         "SettingsLibSettingsTheme",
         "dagger2",
+        "protolog-group",
     ],
     manifest: "quickstep/AndroidManifest.xml",
     min_sdk_version: "current",
@@ -503,7 +567,10 @@
         "Launcher2",
         "Launcher3",
     ],
-    required: ["privapp_whitelist_com.android.launcher3"],
+    required: [
+        "privapp_whitelist_com.android.launcher3",
+        "launcher.quickstep.protolog.pb",
+    ],
 
     resource_dirs: ["quickstep/res"],
 
diff --git a/proguard.flags b/proguard.flags
index 31edd8d..da00c00 100644
--- a/proguard.flags
+++ b/proguard.flags
@@ -57,3 +57,7 @@
 -keep class com.android.quickstep.** {
   *;
 }
+
+-keep class com.android.internal.protolog.** {
+  *;
+}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java b/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java
index 14d391b..bcad478 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java
@@ -86,7 +86,7 @@
         }
         setStateWithAnimationInternal(toState, config, builder);
         builder.addEndListener(success -> {
-            if (!success) {
+            if (!success && !toState.isRecentsViewVisible) {
                 mRecentsView.reset();
             }
         });
diff --git a/quickstep/src/com/android/quickstep/QuickstepProcessInitializer.java b/quickstep/src/com/android/quickstep/QuickstepProcessInitializer.java
index f4e68dc..334bead 100644
--- a/quickstep/src/com/android/quickstep/QuickstepProcessInitializer.java
+++ b/quickstep/src/com/android/quickstep/QuickstepProcessInitializer.java
@@ -25,6 +25,7 @@
 
 import com.android.launcher3.BuildConfig;
 import com.android.launcher3.MainProcessInitializer;
+import com.android.quickstep.util.QuickstepProtoLogGroup;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 
 @SuppressWarnings("unused")
@@ -69,5 +70,7 @@
                     call.descriptor + " called on main thread under " + call.activeTrace
                             + " stackTrace: " + call.stackTrace));
         }
+
+        QuickstepProtoLogGroup.initProtoLog();
     }
 }
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
index 11f1151..83794fb 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
@@ -81,7 +81,7 @@
         // While animating into recents, update the visible task data as needed
         setter.addOnFrameCallback(() -> mRecentsView.loadVisibleTaskData(FLAG_UPDATE_ALL));
         setter.addEndListener(success -> {
-            if (!success) {
+            if (!success && !toState.isRecentsViewVisible()) {
                 mRecentsView.reset();
             }
         });
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index 7f0d665..7f88090 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -19,7 +19,6 @@
 
 import static com.android.quickstep.GestureState.GestureEndTarget.RECENTS;
 import static com.android.quickstep.fallback.RecentsState.DEFAULT;
-import static com.android.quickstep.fallback.RecentsState.HOME;
 import static com.android.quickstep.fallback.RecentsState.MODAL_TASK;
 import static com.android.quickstep.fallback.RecentsState.OVERVIEW_SPLIT_SELECT;
 
@@ -271,7 +270,7 @@
 
     @Override
     public void onStateTransitionComplete(RecentsState finalState) {
-        if (finalState == HOME) {
+        if (!finalState.isRecentsViewVisible()) {
             // Clean-up logic that occurs when recents is no longer in use/visible.
             reset();
         }
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index 73edb9e..8f20cd38 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -18,14 +18,11 @@
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
 
-import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.CLEAR_ALL_BUTTON;
-import static com.android.launcher3.LauncherState.EDIT_MODE;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.LauncherState.OVERVIEW;
 import static com.android.launcher3.LauncherState.OVERVIEW_MODAL_TASK;
 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 android.annotation.TargetApi;
@@ -173,8 +170,7 @@
 
     @Override
     public void onStateTransitionComplete(LauncherState finalState) {
-        if (finalState == NORMAL || finalState == SPRING_LOADED  || finalState == EDIT_MODE
-                || finalState == ALL_APPS) {
+        if (!finalState.isRecentsViewVisible) {
             // Clean-up logic that occurs when recents is no longer in use/visible.
             reset();
         }
diff --git a/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java b/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
new file mode 100644
index 0000000..d0863f8
--- /dev/null
+++ b/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
@@ -0,0 +1,102 @@
+/*
+ * 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.quickstep.util;
+
+import androidx.annotation.NonNull;
+
+import com.android.internal.protolog.ProtoLog;
+import com.android.internal.protolog.common.IProtoLogGroup;
+
+import java.util.UUID;
+
+/** Enums used to interface with the ProtoLog API. */
+public enum QuickstepProtoLogGroup implements IProtoLogGroup {
+
+    ACTIVE_GESTURE_LOG(true, true, false, "ActiveGestureLog");
+
+    private final boolean mEnabled;
+    private volatile boolean mLogToProto;
+    private volatile boolean mLogToLogcat;
+    private final @NonNull String mTag;
+
+    public static void initProtoLog() {
+        ProtoLog.init(QuickstepProtoLogGroup.values());
+    }
+
+    /**
+     * @param enabled     set to false to exclude all log statements for this group from
+     *                    compilation,
+     *                    they will not be available in runtime.
+     * @param logToProto  enable binary logging for the group
+     * @param logToLogcat enable text logging for the group
+     * @param tag         name of the source of the logged message
+     */
+    QuickstepProtoLogGroup(
+            boolean enabled, boolean logToProto, boolean logToLogcat, @NonNull String tag) {
+        this.mEnabled = enabled;
+        this.mLogToProto = logToProto;
+        this.mLogToLogcat = logToLogcat;
+        this.mTag = tag;
+    }
+
+    @Override
+    public boolean isEnabled() {
+        return mEnabled;
+    }
+
+    @Override
+    public boolean isLogToProto() {
+        return mLogToProto;
+    }
+
+    @Override
+    public boolean isLogToLogcat() {
+        return mLogToLogcat;
+    }
+
+    @Override
+    public boolean isLogToAny() {
+        return mLogToLogcat || mLogToProto;
+    }
+
+    @Override
+    public int getId() {
+        return Constants.LOG_START_ID + this.ordinal();
+    }
+
+    @Override
+    public @NonNull String getTag() {
+        return mTag;
+    }
+
+    @Override
+    public void setLogToProto(boolean logToProto) {
+        this.mLogToProto = logToProto;
+    }
+
+    @Override
+    public void setLogToLogcat(boolean logToLogcat) {
+        this.mLogToLogcat = logToLogcat;
+    }
+
+    private static final class Constants {
+
+        private static final int LOG_START_ID =
+                (int) (UUID.nameUUIDFromBytes(QuickstepProtoLogGroup.class.getName().getBytes())
+                        .getMostSignificantBits() % Integer.MAX_VALUE);
+    }
+}
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt
index e4b8069..82a7625 100644
--- a/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt
@@ -97,6 +97,8 @@
     fun bubbleBarView_expanded_threeBubbles() {
         // if we're still expanding, wait with taking a screenshot
         val shouldWait: (ComponentActivity, View) -> Boolean = { _, _ -> bubbleBarView.isExpanding }
+        // increase the frame limit to allow the animation to end before taking the screenshot
+        screenshotRule.frameLimit = 50
         screenshotRule.screenshotTest(
             "bubbleBarView_expanded_threeBubbles",
             checkView = shouldWait,
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 54aea38..5ea7bd9 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -354,7 +354,7 @@
      */
     @Deprecated
     public void reset(Context context) {
-        initGrid(context, getCurrentGridName(context));
+        initGrid(context, getDefaultGridName(context));
     }
 
     @VisibleForTesting
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
index 942b97c..4c017e9 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -17,6 +17,7 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
@@ -129,7 +130,8 @@
             return true;
         }
 
-        if (Flags.enableGridMigrationFix()
+        if (LauncherPrefs.get(context).get(IS_FIRST_LOAD_AFTER_RESTORE)
+                && Flags.enableGridMigrationFix()
                 && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
                 && srcDeviceState.getRows() < destDeviceState.getRows()) {
             Log.i("b/360462379", "Grid migration fix entry point.");
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 21897bf..775d248 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -75,7 +75,6 @@
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LogConfig;
 
-import java.io.File;
 import java.io.InvalidObjectException;
 import java.util.Arrays;
 import java.util.Collection;
@@ -127,12 +126,12 @@
 
         if (Flags.enableNarrowGridRestore()) {
             String oldPhoneFileName = idp.dbFile;
-            List<String> previousDbs = existingDbs();
+            List<String> previousDbs = existingDbs(context);
             removeOldDBs(context, oldPhoneFileName);
             // The idp before this contains data about the old phone, after this it becomes the idp
             // of the current phone.
             idp.reset(context);
-            trySettingPreviousGidAsCurrent(context, idp, oldPhoneFileName, previousDbs);
+            trySettingPreviousGridAsCurrent(context, idp, oldPhoneFileName, previousDbs);
         } else {
             idp.reinitializeAfterRestore(context);
         }
@@ -143,7 +142,7 @@
      * Try setting the gird used in the previous phone to the new one. If the current device doesn't
      * support the previous grid option it will not be set.
      */
-    private static void trySettingPreviousGidAsCurrent(Context context, InvariantDeviceProfile idp,
+    private static void trySettingPreviousGridAsCurrent(Context context, InvariantDeviceProfile idp,
             String oldPhoneDbFileName, List<String> previousDbs) {
         InvariantDeviceProfile.GridOption oldPhoneGridOption = idp.getGridOptionFromFileName(
                 context, oldPhoneDbFileName);
@@ -166,17 +165,19 @@
     /**
      * Returns a list of paths of the existing launcher dbs.
      */
-    private static List<String> existingDbs() {
+    @VisibleForTesting
+    public static List<String> existingDbs(Context context) {
         // At this point idp.dbFile contains the name of the dbFile from the previous phone
         return LauncherFiles.GRID_DB_FILES.stream()
-                .filter(dbName -> new File(dbName).exists())
+                .filter(dbName -> context.getDatabasePath(dbName).exists())
                 .toList();
     }
 
     /**
      * Only keep the last database used on the previous device.
      */
-    private static void removeOldDBs(Context context, String oldPhoneDbFileName) {
+    @VisibleForTesting
+    public static void removeOldDBs(Context context, String oldPhoneDbFileName) {
         // At this point idp.dbFile contains the name of the dbFile from the previous phone
         LauncherFiles.GRID_DB_FILES.stream()
                 .filter(dbName -> !dbName.equals(oldPhoneDbFileName))
diff --git a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
index 479b201..35ac0a1 100644
--- a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
+++ b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
@@ -23,6 +23,7 @@
 import com.android.launcher3.Flags
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.model.ModelDbController
+import com.android.launcher3.provider.RestoreDbTask
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.TestUtil
 import com.android.launcher3.util.rule.BackAndRestoreRule
@@ -67,4 +68,13 @@
             }
         }
     }
+
+    @Test
+    fun testExistingDbsAndRemovingDbs() {
+        var existingDbs = RestoreDbTask.existingDbs(getInstrumentation().targetContext)
+        assert(existingDbs.size == 4)
+        RestoreDbTask.removeOldDBs(getInstrumentation().targetContext, "launcher_4_by_4.db")
+        existingDbs = RestoreDbTask.existingDbs(getInstrumentation().targetContext)
+        assert(existingDbs.size == 1)
+    }
 }