[Gradle] KotlinPluginLifecycle: Ensure no enqueued actions will be missed

^KT-34662 Verification Pending

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt

@@ -385,22 +385,27 @@ private class KotlinPluginLifecycleImpl(override val project: Project) : KotlinP
 
         project.whenEvaluated {
             assert(enqueuedActions.getValue(stage).isEmpty()) { "Expected empty queue from '$stage'" }
-            executeStage(project, stage.nextOrThrow)
+            stage = stage.nextOrThrow
+            executeCurrentStageAndScheduleNext()
         }
     }
 
-    private fun executeStage(project: Project, stage: Stage) {
+    private fun executeCurrentStageAndScheduleNext() {
-        this.stage = stage
+        stage.previousOrNull?.let { previousStage ->
+            assert(enqueuedActions.getValue(previousStage).isEmpty()) {
+                "Actions from previous stage '$previousStage' have not been executed (stage: '$stage')"
+            }
+        }
 
         loopIfNecessary()
 
-        val nextStage = stage.nextOrNull ?: run {
+        stage = stage.nextOrNull ?: run {
             isFinished.set(true)
             return
         }
 
         project.afterEvaluate {
-            executeStage(project, nextStage)
+            executeCurrentStageAndScheduleNext()
         }
     }
 

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt

@@ -350,6 +350,22 @@ class KotlinPluginLifecycleTest {
         }
     }
 
+    @Test
+    fun `test - launching in AfterEvaluate`() = project.runLifecycleAwareTest {
+        val actionInvocations = AtomicInteger(0)
+
+        afterEvaluate {
+            launch {
+                assertEquals(AfterEvaluate.nextOrThrow, currentKotlinPluginLifecycle().stage)
+                assertEquals(1, actionInvocations.incrementAndGet())
+            }
+        }
+
+        await(AfterEvaluate.nextOrThrow.nextOrThrow)
+        assertEquals(1, actionInvocations.get())
+        await(Stage.values.last())
+    }
+
     @Test
     fun `test - Stage - previous next utils`() {
         assertNull(Stage.values.first().previousOrNull)