Fall back to current task when no top-level task is set

lib/annotate_rb/tasks/annotate_models_migrate.rake

@@ -28,8 +28,13 @@ migration_tasks.each do |task|
   next unless Rake::Task.task_defined?(task)
   next if config[:skip_on_db_migrate]
 
-  Rake::Task[task].enhance do # This block is ran after `task` completes
-    task_name = Rake.application.top_level_tasks.last # The name of the task that was run, e.g. "db:migrate"
+  Rake::Task[task].enhance do |current_task| # This block is ran after `task` completes
+    # Prefer the top-level task (the one invoked from the CLI, e.g. "db:migrate") so that
+    # when sub-tasks chain (e.g. db:migrate:redo invokes db:rollback then db:migrate), we
+    # defer the annotation run to after everything completes. Fall back to the currently
+    # enhanced task when there is no top-level task (e.g. when the task is invoked
+    # programmatically from application code rather than from the Rake CLI).
+    task_name = Rake.application.top_level_tasks.last || current_task.name
 
     Rake::Task[task_name].enhance do
       ::AnnotateRb::Runner.run_after_migration

spec/lib/tasks/annotate_models_migrate_spec.rb

@@ -68,6 +68,18 @@
           Rake.application.top_level
         end
       end
+
+      describe "programmatic invocation (no top-level task)" do
+        it "should still annotate model files" do
+          # Simulate programmatic invocation from application code — e.g. Rails'
+          # in-browser "Run pending migrations" button, which calls
+          # Rake::Task["db:migrate"].invoke directly, bypassing Rake's CLI argv parser.
+          Rake.application.instance_variable_set(:@top_level_tasks, [])
+
+          expect(AnnotateRb::Runner).to receive(:run_after_migration)
+          expect { Rake::Task["db:migrate"].invoke }.not_to raise_error
+        end
+      end
     end
 
     context "when skip_on_db_migrate is enabled" do