From 06647e137c4c076e72d7988bfc3772015501c7b7 Mon Sep 17 00:00:00 2001
From: Ian Macartney <366683+ianmacartney@users.noreply.github.com>
Date: Mon, 20 Apr 2026 19:10:09 -0700
Subject: [PATCH] pass in "oneBatchOnly" explicitly to the migration function
 so running from the CLI is always explicit if it wants to start /continue vs.
 do one batch

---
 src/client/index.ts  | 9 +--------
 src/component/lib.ts | 1 +
 src/shared.ts        | 1 +
 3 files changed, 3 insertions(+), 8 deletions(-)

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>;