Add AI docs support for publishing and resolving documentation artifacts

plugin/build.gradle.kts

@@ -38,6 +38,10 @@ gradlePlugin {
         id = "com.automattic.android.publish-to-s3-without-sources"
         implementationClass = "com.automattic.android.publish.PublishToS3WithoutSourcesPlugin"
     }
+    plugins.register("ai-docs") {
+        id = "com.automattic.android.ai-docs"
+        implementationClass = "com.automattic.android.publish.AiDocsPlugin"
+    }
 }
 
 java {

plugin/src/functionalTest/kotlin/com/automattic/android/publish/AiDocsFunctionalTest.kt

@@ -0,0 +1,103 @@
+package com.automattic.android.publish
+
+import org.gradle.testkit.runner.GradleRunner
+import org.gradle.testkit.runner.TaskOutcome
+import java.io.File
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+class AiDocsFunctionalTest {
+    @Test
+    fun `given aiDocs with source directory, when zipAiDocs runs, then produces zip`() {
+        val projectDir = File("build/functionalTest-aiDocs")
+        projectDir.deleteRecursively()
+        projectDir.mkdirs()
+
+        val docsDir = File(projectDir, "docs/ai-reference")
+        docsDir.mkdirs()
+        File(docsDir, "index.md").writeText("# API Reference\n- [posts](posts.md)\n")
+        File(docsDir, "posts.md").writeText("## Posts\n- listWithEditContext(params: PostListParams)\n")
+
+        projectDir.resolve("settings.gradle").writeText("")
+        projectDir.resolve("build.gradle.kts").writeText("""
+            plugins {
+                id("com.automattic.android.publish-to-s3")
+            }
+
+            aiDocs {
+                from(file("docs/ai-reference"))
+            }
+        """.trimIndent())
+
+        val result = GradleRunner.create()
+            .forwardOutput()
+            .withPluginClasspath()
+            .withArguments("zipAiDocs")
+            .withProjectDir(projectDir)
+            .build()
+
+        assertEquals(TaskOutcome.SUCCESS, result.task(":zipAiDocs")?.outcome)
+
+        val zipFile = File(projectDir, "build/ai-docs/ai-docs.zip")
+        assertTrue(zipFile.exists(), "Zip file should exist at ${zipFile.absolutePath}")
+        assertTrue(zipFile.length() > 0, "Zip file should not be empty")
+    }
+
+    @Test
+    fun `given aiDocs consumer plugin with a resolve notation, when listing tasks, then resolveAiDocs is registered`() {
+        val projectDir = File("build/functionalTest-aiDocs-consumer")
+        projectDir.deleteRecursively()
+        projectDir.mkdirs()
+
+        projectDir.resolve("settings.gradle").writeText("")
+        projectDir.resolve("build.gradle.kts").writeText("""
+            plugins {
+                id("com.automattic.android.ai-docs")
+            }
+
+            aiDocs {
+                resolve("com.example:docs:1.0.0")
+            }
+        """.trimIndent())
+
+        val result = GradleRunner.create()
+            .forwardOutput()
+            .withPluginClasspath()
+            .withArguments("tasks", "--all")
+            .withProjectDir(projectDir)
+            .build()
+
+        assertEquals(TaskOutcome.SUCCESS, result.task(":tasks")?.outcome)
+        assertTrue(result.output.contains("resolveAiDocs"), "resolveAiDocs task should be registered")
+    }
+
+    @Test
+    fun `given both publish-to-s3 and ai-docs plugins, when configuring, then the aiDocs extension is shared`() {
+        val projectDir = File("build/functionalTest-aiDocs-both")
+        projectDir.deleteRecursively()
+        projectDir.mkdirs()
+
+        projectDir.resolve("settings.gradle").writeText("")
+        projectDir.resolve("build.gradle.kts").writeText("""
+            plugins {
+                id("com.automattic.android.publish-to-s3")
+                id("com.automattic.android.ai-docs")
+            }
+
+            aiDocs {
+                resolve("com.example:docs:1.0.0")
+            }
+        """.trimIndent())
+
+        val result = GradleRunner.create()
+            .forwardOutput()
+            .withPluginClasspath()
+            .withArguments("tasks", "--all")
+            .withProjectDir(projectDir)
+            .build()
+
+        assertEquals(TaskOutcome.SUCCESS, result.task(":tasks")?.outcome)
+        assertTrue(result.output.contains("resolveAiDocs"), "resolveAiDocs task should be registered")
+    }
+}

plugin/src/main/kotlin/com/automattic/android/publish/AiDocsExtension.kt

@@ -0,0 +1,31 @@
+package com.automattic.android.publish
+
+import org.gradle.api.file.Directory
+import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.model.ObjectFactory
+import org.gradle.api.provider.ListProperty
+import org.gradle.api.provider.Provider
+import java.io.File
+import javax.inject.Inject
+
+open class AiDocsExtension @Inject constructor(objects: ObjectFactory) {
+    val sourceDirectory: DirectoryProperty = objects.directoryProperty()
+
+    val dependencies: ListProperty<String> = objects.listProperty(String::class.java)
+
+    fun from(directory: File) {
+        sourceDirectory.set(directory)
+    }
+
+    fun from(directory: Directory) {
+        sourceDirectory.set(directory)
+    }
+
+    fun from(directory: Provider<Directory>) {
+        sourceDirectory.set(directory)
+    }
+
+    fun resolve(dependencyNotation: String) {
+        dependencies.add(dependencyNotation)
+    }
+}

plugin/src/main/kotlin/com/automattic/android/publish/AiDocsPlugin.kt

@@ -0,0 +1,103 @@
+package com.automattic.android.publish
+
+import org.gradle.api.Plugin
+import org.gradle.api.Project
+import org.gradle.api.publish.PublishingExtension
+import org.gradle.api.publish.maven.MavenPublication
+
+private const val AI_DOCS_CLASSIFIER = "ai-docs"
+private const val NOTATION_WITHOUT_VERSION_PARTS = 2
+private const val NOTATION_WITH_VERSION_PARTS = 3
+
+/**
+ * Consumer-only plugin for resolving AI docs from S3 dependencies.
+ * Does NOT apply maven-publish.
+ */
+class AiDocsPlugin : Plugin<Project> {
+    override fun apply(project: Project) {
+        val extension = project.getOrCreateAiDocsExtension()
+        project.configureAiDocsResolving(extension)
+    }
+}
+
+// Both this consumer plugin and `com.automattic.android.publish-to-s3` register an `aiDocs`
+// extension; reuse the existing one so applying both plugins doesn't fail with a name clash.
+internal fun Project.getOrCreateAiDocsExtension(): AiDocsExtension =
+    extensions.findByType(AiDocsExtension::class.java)
+        ?: extensions.create("aiDocs", AiDocsExtension::class.java)
+
+internal fun Project.configureAiDocsPublishing(extension: AiDocsExtension) {
+    afterEvaluate {
+        if (!extension.sourceDirectory.isPresent) return@afterEvaluate
+
+        val zipTask = tasks.register("zipAiDocs", ZipAiDocsTask::class.java) { task ->
+            task.sourceDirectory.set(extension.sourceDirectory)
+            task.outputZip.set(layout.buildDirectory.file("ai-docs/ai-docs.zip"))
+        }
+
+        extensions.findByType(PublishingExtension::class.java)?.let { publishing ->
+            publishing.publications.withType(MavenPublication::class.java).configureEach { publication ->
+                if (!publication.name.endsWith("PluginMarkerMaven")) {
+                    publication.artifact(zipTask.flatMap { it.outputZip }) {
+                        it.classifier = AI_DOCS_CLASSIFIER
+                        it.extension = "zip"
+                        it.builtBy(zipTask)
+                    }
+                }
+            }
+        }
+    }
+}
+
+internal fun Project.configureAiDocsResolving(extension: AiDocsExtension) {
+    afterEvaluate {
+        val deps = extension.dependencies.getOrElse(emptyList())
+        if (deps.isEmpty()) return@afterEvaluate
+
+        val aiDocsConfig = configurations.maybeCreate("aiDocs").apply {
+            isTransitive = false
+            isCanBeConsumed = false
+        }
+
+        val resolvedNotations = deps.map { notation ->
+            val parts = notation.split(":")
+            require(parts.size == NOTATION_WITHOUT_VERSION_PARTS || parts.size == NOTATION_WITH_VERSION_PARTS) {
+                "Invalid AI docs dependency notation '$notation'; " +
+                    "expected 'group:artifact' or 'group:artifact:version'"
+            }
+            val group = parts[0]
+            val artifact = parts[1]
+            val version = if (parts.size == NOTATION_WITH_VERSION_PARTS) {
+                parts[2]
+            } else {
+                resolveVersionFromDependencyGraph(group, artifact)
+            }
+
+            project.dependencies.add("aiDocs", "$group:$artifact:$version:$AI_DOCS_CLASSIFIER@zip")
+
+            "$group:$artifact:$version"
+        }
+
+        tasks.register("resolveAiDocs", ResolveAiDocsTask::class.java) { task ->
+            task.aiDocsConfiguration = aiDocsConfig
+            task.requestedCoordinates.set(resolvedNotations)
+            // Resolved docs are a build artifact: cleaned by `clean` and implicitly gitignored.
+            // Use this project's build dir (not rootProject) so applying the plugin to multiple
+            // subprojects doesn't race on a shared output directory.
+            task.outputDirectory.set(layout.buildDirectory.dir("ai-docs"))
+        }
+    }
+}
+
+private fun Project.resolveVersionFromDependencyGraph(group: String, artifact: String): String {
+    val matchingDep = configurations
+        .flatMap { it.dependencies }
+        .find { it.group == group && it.name == artifact }
+
+    return matchingDep?.version
+        ?: throw IllegalStateException(
+            "Cannot resolve version for '$group:$artifact'. " +
+                "Either add it as a dependency or specify the version explicitly: " +
+                "resolve(\"$group:$artifact:VERSION\")"
+        )
+}

plugin/src/main/kotlin/com/automattic/android/publish/PublishToS3Plugin.kt

@@ -43,4 +43,7 @@ private fun applyInternal(project: Project, withSources: Boolean) {
     }
 
     project.printPublishedVersionNameAfterPublishTasks()
+
+    val aiDocsExtension = project.getOrCreateAiDocsExtension()
+    project.configureAiDocsPublishing(aiDocsExtension)
 }

plugin/src/main/kotlin/com/automattic/android/publish/ResolveAiDocsTask.kt

@@ -0,0 +1,106 @@
+package com.automattic.android.publish
+
+import org.gradle.api.DefaultTask
+import org.gradle.api.artifacts.Configuration
+import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.provider.ListProperty
+import org.gradle.api.tasks.Input
+import org.gradle.api.tasks.Internal
+import org.gradle.api.tasks.OutputDirectory
+import org.gradle.api.tasks.TaskAction
+import java.io.File
+import java.util.zip.ZipFile
+
+private const val COORDINATE_PARTS = 3
+
+abstract class ResolveAiDocsTask : DefaultTask() {
+    @Internal
+    override fun getDescription() = "Resolves and unpacks AI documentation from dependencies"
+
+    @get:Internal
+    lateinit var aiDocsConfiguration: Configuration
+
+    @get:Input
+    abstract val requestedCoordinates: ListProperty<String>
+
+    @get:OutputDirectory
+    abstract val outputDirectory: DirectoryProperty
+
+    @TaskAction
+    fun resolve() {
+        val outputDir = outputDirectory.get().asFile
+        outputDir.mkdirs()
+
+        val resolvedArtifacts = aiDocsConfiguration.resolvedConfiguration.resolvedArtifacts
+
+        requestedCoordinates.get().forEach { notation ->
+            val parts = notation.split(":")
+            require(parts.size == COORDINATE_PARTS) {
+                "Invalid AI docs coordinate '$notation'; expected 'group:artifact:version'"
+            }
+            val (group, artifact, version) = parts
+
+            val matchingArtifact = resolvedArtifacts.find {
+                it.moduleVersion.id.group == group &&
+                    it.moduleVersion.id.name == artifact &&
+                    it.moduleVersion.id.version == version
+            } ?: run {
+                logger.warn("AI docs artifact not found for $group:$artifact:$version")
+                return@forEach
+            }
+
+            val artifactDir = File(outputDir, "$group/$artifact")
+            val versionDir = File(artifactDir, version)
+
+            cleanStaleVersions(artifactDir, version)
+
+            if (isUpToDate(versionDir)) {
+                logger.lifecycle("AI docs up-to-date: ${versionDir.absolutePath}")
+                return@forEach
+            }
+
+            versionDir.mkdirs()
+            unpackZip(matchingArtifact.file, versionDir)
+            logger.lifecycle("AI docs saved to: ${versionDir.absolutePath}")
+        }
+    }
+
+    private fun isUpToDate(versionDir: File): Boolean =
+        versionDir.exists() && versionDir.list()?.isNotEmpty() == true
+
+    private fun unpackZip(zipFile: File, targetDir: File) {
+        val canonicalTarget = targetDir.canonicalFile.toPath()
+        ZipFile(zipFile).use { zip ->
+            zip.entries().asSequence().forEach { entry ->
+                val targetFile = File(targetDir, entry.name)
+                // Guard against "zip slip": entries must not escape the target directory.
+                require(targetFile.canonicalFile.toPath().startsWith(canonicalTarget)) {
+                    "Zip entry '${entry.name}' would extract outside '$canonicalTarget'"
+                }
+                if (entry.isDirectory) {
+                    targetFile.mkdirs()
+                } else {
+                    unpackEntry(zip, entry, targetFile)
+                }
+            }
+        }
+    }
+
+    private fun unpackEntry(zip: ZipFile, entry: java.util.zip.ZipEntry, targetFile: File) {
+        targetFile.parentFile.mkdirs()
+        zip.getInputStream(entry).use { input ->
+            targetFile.outputStream().use { output -> input.copyTo(output) }
+        }
+    }
+
+    private fun cleanStaleVersions(artifactDir: File, currentVersion: String) {
+        if (!artifactDir.exists()) return
+
+        artifactDir.listFiles()
+            ?.filter { it.isDirectory && it.name != currentVersion }
+            ?.forEach { staleDir ->
+                logger.lifecycle("Removing stale AI docs: ${staleDir.absolutePath}")
+                staleDir.deleteRecursively()
+            }
+    }
+}