Consume all Kotlin sources in KSP to support other code generators

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

diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt
index 3de4417..6fcfabe 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt
@@ -81,14 +81,14 @@
         project: Project,
         kotlinCompilation: KotlinJvmAndroidCompilation,
         kspTaskProvider: TaskProvider<KspAATask>,
-        androidComponent: Component?
+        androidComponent: Component?,
     ) {
         if (androidComponent != null && project.isAgpBuiltInKotlinUsed()) {
             if (project.canUseInternalKspApis()) {
                 val javaSources =
                     (androidComponent.sources.java as? FlatSourceDirectoriesForJavaImpl)?.allButKspAndKaptGenerators()
 
-                val kotlinSources = androidComponent.sources.kotlin?.static
+                val kotlinSources = androidComponent.sources.kotlin?.all
 
                 kspTaskProvider.configure { task ->
                     task.kspConfig.javaSourceRoots.from(javaSources)
@@ -242,7 +242,12 @@
         // Order is important here as we update task with AGP generated sources and
         // then update AGP with source that KSP will generate.
         // Mixing this up will cause circular dependency in Gradle
-        tryUpdateKspWithAndroidSourceSets(project, kotlinCompilation, kspTaskProvider, androidComponent)
+        tryUpdateKspWithAndroidSourceSets(
+            project,
+            kotlinCompilation,
+            kspTaskProvider,
+            androidComponent,
+        )
 
         registerGeneratedSources(
             project,

integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/AndroidBuiltInKotlinIT.kt

diff --git a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/AndroidBuiltInKotlinIT.kt b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/AndroidBuiltInKotlinIT.kt
index c091f29..beec8f5 100644
--- a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/AndroidBuiltInKotlinIT.kt
+++ b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/AndroidBuiltInKotlinIT.kt
@@ -289,4 +289,59 @@
             }
         }
     }
+
+    @Test
+    fun `test Kotlin generated sources are picked up by KSP with built-in Kotlin`() {
+        val gradleRunner = GradleRunner.create().withProjectDir(project.root).forwardOutput()
+
+        File(project.root, "workload/build.gradle.kts").appendText(
+            """
+                android {
+                    androidComponents.onVariants { variant ->
+                        val name = variant.name
+                
+                        val task = project.tasks.register("generate${'$'}{name.replaceFirstChar(Char::uppercase)}KotlinSource", KotlinSourceGenerationTask::class) {
+                            getOutputDirectory().set(project.layout.buildDirectory.dir("generated/custom_kotlin/${'$'}{name}"))
+                        }
+                
+                        variant.sources.kotlin?.addGeneratedSourceDirectory(task, KotlinSourceGenerationTask::getOutputDirectory)
+                    }
+                }
+
+                abstract class KotlinSourceGenerationTask : DefaultTask() {
+                    @OutputDirectory
+                    abstract fun getOutputDirectory(): DirectoryProperty
+                    @TaskAction
+                    fun generateCode() {
+                        val outDir = getOutputDirectory().get().asFile
+                        val outFile = File(outDir, "com/example/GeneratedKotlinClass.kt")
+                        outFile.parentFile.mkdirs()
+                        outFile.writeText(""${'"'}
+                            package com.example
+                            class GeneratedKotlinClass {
+                                fun foo(): String = "Generated"
+                            }
+                        ""${'"'}.trimIndent())
+                    }
+                }
+            """.trimIndent()
+        )
+
+        File(project.root, "gradle.properties").appendText("\nagpVersion=9.0.0-beta05")
+
+        gradleRunner.withArguments(":workload:assembleDebug", "--stacktrace").build().let { result ->
+            Assert.assertEquals(
+                TaskOutcome.SUCCESS, result.task(":workload:kspDebugKotlin")?.outcome
+            )
+
+            val logFile = File(
+                project.root, "workload/build/generated/ksp/debug/resources/TestProcessor.log"
+            )
+            Assert.assertTrue(
+                "Log file expected to exist but it does not exist: ${logFile.path}", logFile.exists()
+            )
+            val logContent = logFile.readText()
+            Assert.assertTrue(logContent.contains("Found GeneratedKotlinClass: com.example.GeneratedKotlinClass"))
+        }
+    }
 }

integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt

diff --git a/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt b/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
index 57d4798..aaf015c 100644
--- a/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
+++ b/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
@@ -58,6 +58,9 @@
                 logger.warn("[$moduleName] Mangled name for internalFun: $internalName")
             }
         }
+        resolver.getClassDeclarationByName("com.example.GeneratedKotlinClass")?.let {
+            emit("Found GeneratedKotlinClass: ${it.qualifiedName?.asString()}", "")
+        }
         invoked = true
         return emptyList()
     }