Allow enabling KSP2 per module

gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspExtension.kt

@@ -17,14 +17,24 @@
 
 package com.google.devtools.ksp.gradle
 
+import com.google.devtools.ksp.KspExperimental
 import org.gradle.api.GradleException
 import org.gradle.api.Project
 import org.gradle.api.file.ConfigurableFileCollection
+import org.gradle.api.provider.Property
 import org.gradle.api.provider.Provider
 import org.gradle.process.CommandLineArgumentProvider
 import javax.inject.Inject
 
 abstract class KspExtension @Inject constructor(project: Project) {
+    /**
+     * Enables or disables KSP 2, defaults to the `ksp.useKsp2` gradle property or `false` if that's not set.
+     *
+     * This API is temporary and will be removed once KSP1 is removed.
+     */
+    @KspExperimental
+    abstract val useKsp2: Property<Boolean>
+
     internal val apOptions = project.objects.mapProperty(String::class.java, String::class.java)
     internal val commandLineArgumentProviders = project.objects.listProperty(CommandLineArgumentProvider::class.java)
         .also { it.finalizeValueOnRead() }

gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt

@@ -16,6 +16,7 @@
  */
 package com.google.devtools.ksp.gradle
 
+import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.gradle.model.builder.KspModelBuilder
 import org.gradle.api.Action
 import org.gradle.api.Project
@@ -56,15 +57,25 @@ import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
 import org.jetbrains.kotlin.gradle.plugin.SubpluginArtifact
 import org.jetbrains.kotlin.gradle.plugin.SubpluginOption
 import org.jetbrains.kotlin.gradle.plugin.getKotlinPluginVersion
-import org.jetbrains.kotlin.gradle.plugin.mpp.*
-import org.jetbrains.kotlin.gradle.tasks.*
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinCommonCompilation
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmAndroidCompilation
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmCompilation
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinSharedNativeCompilation
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinWithJavaCompilation
+import org.jetbrains.kotlin.gradle.tasks.AbstractKotlinCompile
+import org.jetbrains.kotlin.gradle.tasks.AbstractKotlinCompileTool
+import org.jetbrains.kotlin.gradle.tasks.BaseKotlinCompile
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompilationTask
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+import org.jetbrains.kotlin.gradle.tasks.KotlinNativeCompile
 import org.jetbrains.kotlin.incremental.isJavaFile
 import org.jetbrains.kotlin.incremental.isKotlinFile
 import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
 import java.io.File
 import java.util.concurrent.Callable
 import javax.inject.Inject
 
+@OptIn(KspExperimental::class)
 class KspGradleSubplugin @Inject internal constructor(private val registry: ToolingModelBuilderRegistry) :
     KotlinCompilerPluginSupportPlugin {
     companion object {
@@ -197,7 +208,13 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
     private lateinit var kspConfigurations: KspConfigurations
 
     override fun apply(target: Project) {
-        target.extensions.create("ksp", KspExtension::class.java)
+        val ksp = target.extensions.create("ksp", KspExtension::class.java)
+        ksp.useKsp2.convention(
+            target.providers
+                .gradleProperty("ksp.useKSP2")
+                .map { it.toBoolean() }
+                .orElse(false)
+        )
         kspConfigurations = KspConfigurations(target)
         registry.register(KspModelBuilder())
     }
@@ -430,7 +447,9 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
         val isIntermoduleIncremental =
             (project.providers.gradleProperty("ksp.incremental.intermodule").orNull?.toBoolean() ?: true) &&
                 isIncremental
-        val useKSP2 = project.providers.gradleProperty("ksp.useKSP2").orNull?.toBoolean() ?: false
+        val useKSP2 = kspExtension.useKsp2
+            .apply { finalizeValue() }
+            .get()
 
         // Create and configure KSP tasks.
         @Suppress("DEPRECATION") val kspTaskProvider = if (useKSP2) {

gradle-plugin/src/test/kotlin/com/google/devtools/ksp/gradle/GradleCompilationTest.kt

@@ -387,4 +387,17 @@ class GradleCompilationTest {
         assertThat(result.output).contains("HAS LIBRARY: ")
         assertThat(result.output).doesNotContain("app/build/generated/ksp/main/classes")
     }
+
+    @Test
+    fun changingKsp2AtRuntime() {
+        testRule.setupAppAsJvmApp()
+        testRule.appModule.buildFileAdditions.add(
+            """
+                @OptIn(com.google.devtools.ksp.KspExperimental::class)
+                ksp { useKsp2.set(true) }
+            """.trimIndent()
+        )
+
+        testRule.runner().withArguments().build()
+    }
 }