diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index 211c351..0732379 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -120,7 +120,8 @@
             @NonNull DeviceGridState destDeviceState,
             @NonNull DatabaseHelper target,
             @NonNull SQLiteDatabase source,
-            boolean isDestNewDb) {
+            boolean isDestNewDb,
+            ModelDelegate modelDelegate) {
 
         if (!needsToMigrate(srcDeviceState, destDeviceState)) {
             return true;
@@ -174,7 +175,6 @@
             return true;
         } catch (Exception e) {
             Log.e(TAG, "Error during grid migration", e);
-
             return false;
         } finally {
             Log.v(TAG, "Workspace migration completed in "
@@ -182,6 +182,8 @@
 
             // Save current configuration, so that the migration does not run again.
             destDeviceState.writeToPrefs(context);
+            // Notify if we've migrated successfully
+            modelDelegate.gridMigrationComplete(srcDeviceState, destDeviceState);
         }
     }
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
index 0b12af8..1729153 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -53,6 +53,7 @@
         target: DatabaseHelper,
         source: SQLiteDatabase,
         isDestNewDb: Boolean,
+        modelDelegate: ModelDelegate,
     ) {
         if (!GridSizeMigrationDBController.needsToMigrate(srcDeviceState, destDeviceState)) {
             return
@@ -132,6 +133,9 @@
 
             // Save current configuration, so that the migration does not run again.
             destDeviceState.writeToPrefs(context)
+
+            // Notify if we've migrated successfully
+            modelDelegate.gridMigrationComplete(srcDeviceState, destDeviceState)
         }
     }
 
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 4e57944..44b7e8b 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -438,12 +438,12 @@
         ModelDbController dbController = mApp.getModel().getModelDbController();
         if (Flags.gridMigrationRefactor()) {
             try {
-                dbController.attemptMigrateDb(restoreEventLogger);
+                dbController.attemptMigrateDb(restoreEventLogger, mModelDelegate);
             } catch (Exception e) {
                 FileLog.e(TAG, "Failed to migrate grid", e);
             }
         } else {
-            dbController.tryMigrateDB(restoreEventLogger);
+            dbController.tryMigrateDB(restoreEventLogger, mModelDelegate);
         }
         Log.d(TAG, "loadWorkspace: loading default favorites if necessary");
         dbController.loadDefaultFavoritesIfNecessary();
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index e76391f..0138390 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -361,7 +361,8 @@
     /**
      * Migrates the DB. If the migration failed, it clears the DB.
      */
-    public void attemptMigrateDb(LauncherRestoreEventLogger restoreEventLogger) throws Exception {
+    public void attemptMigrateDb(LauncherRestoreEventLogger restoreEventLogger,
+            ModelDelegate modelDelegate) throws Exception {
         createDbIfNotExists();
 
         if (shouldResetDb()) {
@@ -389,7 +390,7 @@
             boolean isDestNewDb = !existingDBs.contains(destDeviceState.getDbFile());
             GridSizeMigrationLogic gridSizeMigrationLogic = new GridSizeMigrationLogic();
             gridSizeMigrationLogic.migrateGrid(mContext, srcDeviceState, destDeviceState,
-                    mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb);
+                    mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb, modelDelegate);
         } catch (Exception e) {
             resetLauncherDb(restoreEventLogger);
             throw new Exception("attemptMigrateDb: Failed to migrate grid", e);
@@ -403,8 +404,9 @@
     /**
      * Migrates the DB if needed. If the migration failed, it clears the DB.
      */
-    public void tryMigrateDB(@Nullable LauncherRestoreEventLogger restoreEventLogger) {
-        if (!migrateGridIfNeeded()) {
+    public void tryMigrateDB(@Nullable LauncherRestoreEventLogger restoreEventLogger,
+            ModelDelegate modelDelegate) {
+        if (!migrateGridIfNeeded(modelDelegate)) {
             if (restoreEventLogger != null) {
                 if (LauncherPrefs.get(mContext).get(NO_DB_FILES_RESTORED)) {
                     restoreEventLogger.logLauncherItemsRestoreFailed(DATA_TYPE_DB_FILE, 1,
@@ -434,7 +436,7 @@
      * @return true if migration was success or ignored, false if migration failed
      * and the DB should be reset.
      */
-    private boolean migrateGridIfNeeded() {
+    private boolean migrateGridIfNeeded(ModelDelegate modelDelegate) {
         createDbIfNotExists();
         if (LauncherPrefs.get(mContext).get(getEmptyDbCreatedKey())) {
             // If we have already create a new DB, ignore migration
@@ -468,7 +470,8 @@
             DeviceGridState destDeviceState = new DeviceGridState(idp);
             boolean isDestNewDb = !existingDBs.contains(destDeviceState.getDbFile());
             return GridSizeMigrationDBController.migrateGridIfNeeded(mContext, srcDeviceState,
-                    destDeviceState, mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb);
+                    destDeviceState, mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb,
+                    modelDelegate);
         } catch (Exception e) {
             FileLog.e(TAG, "migrateGridIfNeeded: Failed to migrate grid", e);
             return false;
diff --git a/src/com/android/launcher3/model/ModelDelegate.java b/src/com/android/launcher3/model/ModelDelegate.java
index 2264d35..5a2aef0 100644
--- a/src/com/android/launcher3/model/ModelDelegate.java
+++ b/src/com/android/launcher3/model/ModelDelegate.java
@@ -123,6 +123,11 @@
     @WorkerThread
     public void modelLoadComplete() { }
 
+    /** Called when grid migration has completed as part of grid size refactor. */
+    @WorkerThread
+    public void gridMigrationComplete(
+            @NonNull DeviceGridState src, @NonNull DeviceGridState dest) { }
+
     /**
      * Called when the delegate is no loner needed
      */
diff --git a/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java b/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
index a9082e2..47110fa 100644
--- a/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
+++ b/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
@@ -61,9 +61,11 @@
             ModelDbController controller = model.getModelDbController();
             // Migrate any previous data so that the DB state is correct
             if (Flags.gridMigrationRefactor()) {
-                controller.attemptMigrateDb(null /* restoreEventLogger */);
+                controller.attemptMigrateDb(
+                        null /* restoreEventLogger */, model.getModelDelegate());
             } else {
-                controller.tryMigrateDB(null /* restoreEventLogger */);
+                controller.tryMigrateDB(null /* restoreEventLogger */,
+                        model.getModelDelegate());
             }
 
             // Create DB again to load fresh data
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt b/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
index 8d072d8..547d05e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
@@ -31,8 +31,9 @@
         loadModelSync()
         TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
             modelDbController.run {
-                if (Flags.gridMigrationRefactor()) attemptMigrateDb(null /* restoreEventLogger */)
-                else tryMigrateDB(null /* restoreEventLogger */)
+                if (Flags.gridMigrationRefactor())
+                    attemptMigrateDb(null /* restoreEventLogger */, modelDelegate)
+                else tryMigrateDB(null /* restoreEventLogger */, modelDelegate)
                 createEmptyDB()
                 clearEmptyDbFlag()
             }
@@ -74,7 +75,7 @@
         loadModelSync()
         TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
             val controller: ModelDbController = modelDbController
-            controller.attemptMigrateDb(null /* restoreEventLogger */)
+            controller.attemptMigrateDb(null /* restoreEventLogger */, modelDelegate)
             modelDbController.newTransaction().use { transaction ->
                 val values =
                     ContentValues().apply {
diff --git a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
index b4ee090..38fad6b 100644
--- a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
+++ b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
@@ -25,6 +25,7 @@
 import com.android.launcher3.Flags
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.model.ModelDbController
+import com.android.launcher3.model.ModelDelegate
 import com.android.launcher3.provider.RestoreDbTask
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.TestUtil
@@ -34,6 +35,7 @@
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
 
 /**
  * Makes sure to test {@code RestoreDbTask#removeOldDBs}, we need to remove all the dbs that are not
@@ -49,6 +51,8 @@
     @Rule
     val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
 
+    val modelDelegate = mock<ModelDelegate>()
+
     @Before
     fun setUp() {
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_NARROW_GRID_RESTORE)
@@ -68,9 +72,9 @@
     fun oldDatabasesNotPresentAfterRestore() {
         val dbController = ModelDbController(getInstrumentation().targetContext)
         if (Flags.gridMigrationRefactor()) {
-            dbController.attemptMigrateDb(null)
+            dbController.attemptMigrateDb(null, modelDelegate)
         } else {
-            dbController.tryMigrateDB(null)
+            dbController.tryMigrateDB(null, modelDelegate)
         }
         TestUtil.runOnExecutorSync(MODEL_EXECUTOR) {
             assert(backAndRestoreRule.getDatabaseFiles().size == 1) {
diff --git a/tests/src/com/android/launcher3/model/GridMigrationTest.kt b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
index b8ffe74..8bd0c60 100644
--- a/tests/src/com/android/launcher3/model/GridMigrationTest.kt
+++ b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
@@ -32,6 +32,8 @@
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
 
 private val phoneContext = InstrumentationRegistry.getInstrumentation().targetContext
 
@@ -87,6 +89,8 @@
     @Rule
     val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
 
+    val modelDelegate = mock<ModelDelegate>()
+
     @Before
     fun setup() {
         setFlagsRule.setFlags(true, Flags.FLAG_ONE_GRID_SPECS)
@@ -102,6 +106,7 @@
                 dst.dbHelper,
                 src.dbHelper.readableDatabase,
                 true,
+                modelDelegate,
             )
         } else {
             GridSizeMigrationDBController.migrateGridIfNeeded(
@@ -111,6 +116,7 @@
                 dst.dbHelper,
                 src.dbHelper.readableDatabase,
                 true,
+                modelDelegate,
             )
         }
     }
@@ -149,11 +155,12 @@
     }
 
     /**
-     * Migrate src into dst and compare to target. This method validates 3 things:
+     * Migrate src into dst and compare to target. This method validates 4 things:
      * 1. dst has the same number of items as src after the migration, meaning, none of the items
      *    were removed during the migration.
      * 2. dst is valid, meaning that none of the items overlap with each other.
      * 3. dst is equal to target to ensure we don't unintentionally change the migration logic.
+     * 4. migration notifies the complete callback.
      */
     private fun runTest(src: GridMigrationData, dst: GridMigrationData, target: GridMigrationData) {
         migrate(src, dst)
@@ -162,6 +169,7 @@
         }
         validateDb(dst)
         compare(dst, target, src)
+        verify(modelDelegate).gridMigrationComplete(src.gridState, dst.gridState)
     }
 
     // Copying the src db for all tests.