diff --git a/includes/class-migration.php b/includes/class-migration.php
index cae22ba0e..dac8e98bd 100644
--- a/includes/class-migration.php
+++ b/includes/class-migration.php
@@ -23,7 +23,7 @@ class Migration {
 	 */
 	public static function init() {
 		\add_action( 'activitypub_migrate', array( self::class, 'async_migration' ) );
-		\add_action( 'activitypub_upgrade', array( self::class, 'async_upgrade' ) );
+		\add_action( 'activitypub_upgrade', array( self::class, 'async_upgrade' ), 10, 99 );
 		\add_action( 'activitypub_update_comment_counts', array( self::class, 'update_comment_counts' ), 10, 2 );
 
 		self::maybe_migrate();
@@ -221,22 +221,29 @@ public static function async_migration( $version_from_db ) {
 	 * @param mixed  ...$args  The arguments to pass to the callback.
 	 */
 	public static function async_upgrade( $callback, ...$args ) {
+		$args = \func_get_args();
+
 		// Bail if the existing lock is still valid.
 		if ( self::is_locked() ) {
-			\wp_schedule_single_event( time() + MINUTE_IN_SECONDS, 'activitypub_upgrade', \array_merge( array( $callback ), $args ) );
+			\wp_schedule_single_event( time() + MINUTE_IN_SECONDS, 'activitypub_upgrade', $args );
 			return;
 		}
 
 		self::lock();
 
-		$next = \call_user_func_array( array( self::class, $callback ), $args );
+		$callback = array_shift( $args );  // Remove $callback from arguments.
+		$next     = \call_user_func_array( array( self::class, $callback ), $args );
+
+		self::unlock();
 
 		if ( ! empty( $next ) ) {
 			// Schedule the next run, adding the result to the arguments.
-			\wp_schedule_single_event( \time() + 30, 'activitypub_upgrade', \array_merge( array( $callback ), $next ) );
+			\wp_schedule_single_event(
+				\time() + 30,
+				'activitypub_upgrade',
+				\array_merge( array( $callback ), \array_values( $next ) )
+			);
 		}
-
-		self::unlock();
 	}
 
 	/**
diff --git a/tests/includes/class-test-migration.php b/tests/includes/class-test-migration.php
index 7a27de9e0..77afe6d47 100644
--- a/tests/includes/class-test-migration.php
+++ b/tests/includes/class-test-migration.php
@@ -534,27 +534,13 @@ public function test_async_upgrade() {
 
 		// Test scheduling next batch when callback returns more work.
 		Migration::async_upgrade( 'create_post_outbox_items', 1, 0 ); // Small batch size to force multiple batches.
-		$scheduled = \wp_next_scheduled(
-			'activitypub_upgrade',
-			array(
-				'create_post_outbox_items',
-				'batch_size' => 1,
-				'offset'     => 1,
-			)
-		);
+		$scheduled = \wp_next_scheduled( 'activitypub_upgrade', array( 'create_post_outbox_items', 1, 1 ) );
 		$this->assertNotFalse( $scheduled );
 
 		// Test no scheduling when callback returns null (no more work).
 		Migration::async_upgrade( 'create_post_outbox_items', 100, 1000 ); // Large offset to ensure no posts found.
 		$this->assertFalse(
-			\wp_next_scheduled(
-				'activitypub_upgrade',
-				array(
-					'create_post_outbox_items',
-					'batch_size' => 100,
-					'offset'     => 1100,
-				)
-			)
+			\wp_next_scheduled( 'activitypub_upgrade', array( 'create_post_outbox_items', 100, 1100 ) )
 		);
 	}
 
@@ -566,16 +552,7 @@ public function test_async_upgrade() {
 	public function test_async_upgrade_multiple_args() {
 		// Test that multiple arguments are passed correctly.
 		Migration::async_upgrade( 'update_comment_counts', 50, 100 );
-		$scheduled = \wp_next_scheduled(
-			'activitypub_upgrade',
-			array(
-				'update_comment_counts',
-				array(
-					'batch_size' => 50,
-					'offset'     => 150,
-				),
-			)
-		);
+		$scheduled = \wp_next_scheduled( 'activitypub_upgrade', array( 'update_comment_counts', 50, 150 ) );
 		$this->assertFalse( $scheduled, 'Should not schedule next batch when no comments found' );
 	}