From 5223fc01d4f256fdc31db202c171a0ec2f561398 Mon Sep 17 00:00:00 2001
From: Adwait Kumar Singh <adwsingh@amazon.com>
Date: Mon, 11 Aug 2025 16:58:18 -0700
Subject: [PATCH] Fix service file merging for mcp-schemas

---
 mcp/mcp-schemas/build.gradle.kts | 102 +++++++++++++++----------------
 1 file changed, 51 insertions(+), 51 deletions(-)

diff --git a/mcp/mcp-schemas/build.gradle.kts b/mcp/mcp-schemas/build.gradle.kts
index 60b760634..e49ad08ae 100644
--- a/mcp/mcp-schemas/build.gradle.kts
+++ b/mcp/mcp-schemas/build.gradle.kts
@@ -39,6 +39,7 @@ afterEvaluate {
                 srcDir(typePath)
                 srcDir(serverPath)
                 include("META-INF/**")
+                exclude("META-INF/services/**") // Exclude original service files, use merged ones instead
             }
         }
     }
@@ -48,62 +49,61 @@ tasks.named("compileJava") {
     dependsOn("smithyBuild")
 }
 
-// Needed because sources-jar needs to run after smithy-build is done
-tasks.sourcesJar {
-    mustRunAfter(tasks.compileJava)
-    configureServiceFileMerging()
-}
-
-tasks.processResources {
-    dependsOn(tasks.compileJava)
-    configureServiceFileMerging()
-}
-
-interface Injected {
-    @get:Inject val fs: FileSystemOperations
-}
-
-fun AbstractCopyTask.configureServiceFileMerging() {
-    val serviceEntries = mutableMapOf<String, MutableSet<String>>()
-    val tempServicesDir = temporaryDir // Capture at configuration time
-
-    // Configure Jar tasks to include temp directory at configuration time
-    if (this is Jar) {
-        from(tempServicesDir)
-    }
-
-    eachFile {
-        if (path.startsWith("META-INF/services/")) {
-            val serviceName = path.substring("META-INF/services/".length)
-
-            if (!serviceEntries.containsKey(serviceName)) {
-                serviceEntries[serviceName] = mutableSetOf()
+// TODO remove once we move to codegen modes instead of plugins.
+val serviceFilesMerger =
+    tasks.register("mergeServiceFiles") {
+        dependsOn(tasks.smithyBuild)
+
+        val outputServiceDir = layout.buildDirectory.dir("merged-services/META-INF/services")
+        outputs.dir(outputServiceDir)
+
+        val projectDir = project.projectDir
+
+        doLast {
+            // Use hardcoded paths because of https://docs.gradle.org/8.14.3/userguide/configuration_cache.html#config_cache:requirements:disallowed_types
+            val sourceServiceDirs =
+                listOf(
+                    File(projectDir, "build/smithyprojections/mcp-schemas/source/java-type-codegen/META-INF/services"),
+                    File(projectDir, "build/smithyprojections/mcp-schemas/source/java-server-codegen/META-INF/services"),
+                )
+
+            val serviceEntries = mutableMapOf<String, MutableSet<String>>()
+
+            sourceServiceDirs.forEach { serviceDir ->
+                if (serviceDir.exists() && serviceDir.isDirectory) {
+                    serviceDir.listFiles()?.forEach { serviceFile ->
+                        if (serviceFile.isFile) {
+                            val serviceName = serviceFile.name
+                            serviceEntries
+                                .computeIfAbsent(serviceName) { mutableSetOf() }
+                                .addAll(serviceFile.readLines().map { it.trim() })
+                        }
+                    }
+                }
             }
 
-            serviceEntries[serviceName]!!.addAll(
-                file
-                    .readLines()
-                    .filter { it.trim().isNotEmpty() }
-                    .map { it.trim() },
-            )
+            val outputDir = outputServiceDir.get().asFile
+            outputDir.mkdirs()
 
-            exclude()
+            serviceEntries.forEach { (serviceName, lines) ->
+                val serviceFile = File(outputDir, serviceName)
+                serviceFile.writeText(lines.sorted().joinToString("\n") + "\n")
+            }
         }
     }
 
-    doLast {
-        val outputDir =
-            when (this) {
-                is ProcessResources -> this.destinationDir
-                else -> tempServicesDir
-            }
-
-        val servicesDir = File(outputDir, "META-INF/services")
-        servicesDir.mkdirs()
-
-        serviceEntries.forEach { (serviceName, lines) ->
-            val serviceFile = File(servicesDir, serviceName)
-            serviceFile.writeText(lines.joinToString("\n") + "\n")
-        }
+// processResources will include merged service files in the main resources
+tasks.processResources {
+    dependsOn(serviceFilesMerger)
+    from(layout.buildDirectory.dir("merged-services")) {
+        into(".")
     }
+    duplicatesStrategy = DuplicatesStrategy.EXCLUDE
+}
+
+// Ensure sourcesJar waits for smithyBuild to complete and includes merged service files
+tasks.sourcesJar {
+    dependsOn(tasks.smithyBuild, serviceFilesMerger)
+    from(layout.buildDirectory.dir("merged-services"))
+    duplicatesStrategy = DuplicatesStrategy.EXCLUDE
 }