diff --git a/src/client/index.ts b/src/client/index.ts
index a9cf6e5..549ded7 100644
--- a/src/client/index.ts
+++ b/src/client/index.ts
@@ -294,14 +294,7 @@ export class Migrations<DataModel extends GenericDataModel> {
         // Auto-detect function name for direct CLI/dashboard invocation.
         // The component always provides cursor and dryRun when scheduling
         // batches, so if either is missing this is a direct invocation.
-        if (
-          args.fn ||
-          args.next ||
-          args.cursor === undefined ||
-          args.cursor === "" ||
-          args.dryRun === undefined ||
-          args.batchSize === 0
-        ) {
+        if (!args.oneBatchOnly) {
           if (args.cursor === "" || args.batchSize === 0) {
             console.warn(
               "Running this from the CLI or dashboard? Here's some args to use:",
diff --git a/src/component/lib.ts b/src/component/lib.ts
index 2aadb69..766297c 100644
--- a/src/component/lib.ts
+++ b/src/component/lib.ts
@@ -117,6 +117,7 @@ export const migrate = mutation({
             cursor: state.cursor,
             batchSize,
             dryRun,
+            oneBatchOnly: true,
           },
         );
         updateState(result);
diff --git a/src/shared.ts b/src/shared.ts
index 4c4961c..ccd609e 100644
--- a/src/shared.ts
+++ b/src/shared.ts
@@ -7,6 +7,7 @@ export const migrationArgs = {
   dryRun: v.optional(v.boolean()),
   next: v.optional(v.array(v.string())),
   reset: v.optional(v.boolean()),
+  oneBatchOnly: v.optional(v.boolean()),
 };
 export type MigrationArgs = ObjectType<typeof migrationArgs>;