fix: align tool validation and source updates with MCP SDK

2026-08-06 00:48:50 +09:00
parent 4a343d018a
commit a88122bd09
11 changed files with 255 additions and 139 deletions
--- a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/annotation/ToolHint.java
+++ b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/annotation/ToolHint.java
@@ -25,6 +25,7 @@ import java.lang.annotation.Target;
 public @interface ToolHint {
    boolean register() default false;
    boolean requiresApproval() default false;
+    String group() default "";
    String mappingId() default "";
    String inputSchemaResource() default "";
    String outputSchemaResource() default "";
--- a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/ToolRegistryHeartbeatSender.java
+++ b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/ToolRegistryHeartbeatSender.java
@@ -107,7 +107,8 @@ public class ToolRegistryHeartbeatSender {
                    meta.setTimeoutMillis(5000L);
                    meta.setEnabled(true);
                    meta.setDescription(functionAnnotation.description());
-                    meta.setCategoryKey("default");
+                    meta.setCategoryKey(hintAnnotation == null || hintAnnotation.group().isBlank()
+                            ? "default" : hintAnnotation.group());
                    meta.setIntegrationType("REST");
                    meta.setMciServiceId(hintAnnotation != null ? hintAnnotation.mappingId() : "");
                    meta.setPodUrl(podUrl);
--- a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/util/ToolSourceUpdater.java
+++ b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/util/ToolSourceUpdater.java
@@ -1,123 +1,140 @@
 package io.shinhanlife.dap.lib.util;

-
-/**
- * @package io.shinhanlife.dap.lib.util
- * @className ToolSourceUpdater
- * @description AX HUB 시스템 처리 클래스
- * @author 0986406
- * @create 2026.09.01
- * <pre>
- * ---------- 개정이력 ----------
- * 수정일      수정자    수정내용
- * ---------- -------- ---------------------------
- * 2026.09.01  0986406    최초생성
- * 
- * </pre>
- */
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.List;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
+import java.util.Comparator;

-public class ToolSourceUpdater {
+/** Updates MCP SDK and project-owned metadata in a generated tool source file. */
+public final class ToolSourceUpdater {

-    public static void updateToolSource(String toolName, String domainGroup, String description, boolean register, Boolean requiresApproval) throws Exception {
-        // 1. Find all *UseCase.java files in dap-was-* directories
-        String envSourceDir = System.getenv("AXHUB_SOURCE_DIR");
-        Path rootDir = envSourceDir != null ? Paths.get(envSourceDir) : Paths.get(".");
-        
-        List<Path> javaFiles;
-        try (Stream<Path> paths = Files.walk(rootDir)) {
-            javaFiles = paths
-                .filter(Files::isRegularFile)
-                .filter(p -> p.toString().endsWith("UseCase.java"))
-                .filter(p -> p.toString().contains("dap-was-") || p.toString().contains("axhub-tool-"))
-                .collect(Collectors.toList());
-        }
+    private ToolSourceUpdater() {
+    }

-        Path targetFile = null;
-        String content = null;
-
-        // 2. Find the specific file for the tool
-        String functionName = toolName;
-        if (toolName.contains("_")) {
-            functionName = toolName.substring(toolName.indexOf("_") + 1);
-        }
-
-        Pattern namePattern = Pattern.compile("@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(toolName) + "\"", Pattern.DOTALL);
-        Pattern namePattern2 = Pattern.compile("@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(functionName) + "\"", Pattern.DOTALL);
-
-        for (Path path : javaFiles) {
-            String text = Files.readString(path);
-            if (namePattern.matcher(text).find()) {
-                targetFile = path;
-                content = text;
-                break;
-            } else if (namePattern2.matcher(text).find()) {
-                targetFile = path;
-                content = text;
-                toolName = functionName; // Use baseName for subsequent replacements
-                break;
-            }
-        }
+    public static void updateToolSource(String toolName, String domainGroup, String description,
+                                        boolean register, Boolean requiresApproval) throws IOException {
+        String configuredSourceDirectory = System.getenv("AXHUB_SOURCE_DIR");
+        Path rootDirectory = configuredSourceDirectory == null || configuredSourceDirectory.isBlank()
+                ? Paths.get(".") : Paths.get(configuredSourceDirectory);
+        updateToolSource(rootDirectory, toolName, domainGroup, description, register, requiresApproval);
+    }

+    static void updateToolSource(Path rootDirectory, String toolName, String domainGroup, String description,
+                                 boolean register, Boolean requiresApproval) throws IOException {
+        Path targetFile = findToolSource(rootDirectory, toolName);
        if (targetFile == null) {
-            throw new Exception("소스 코드를 찾을 수 없습니다: " + toolName);
+            throw new IllegalArgumentException("Tool source not found: " + toolName);
        }

-        // 3. Update @McpTool group
-        if (domainGroup != null && !domainGroup.trim().isEmpty()) {
-            Pattern groupPattern = Pattern.compile("(@McpTool\\s*\\([^)]*group\\s*=\\s*\")([^\"]+)(\")", Pattern.DOTALL);
-            Matcher groupMatcher = groupPattern.matcher(content);
-            if (groupMatcher.find()) {
-                content = groupMatcher.replaceFirst("$1" + domainGroup + "$3");
-            }
-        }
-
-        // 4. Update @McpFunction description
-        if (description != null) {
-            Pattern funcPattern = Pattern.compile("(@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(toolName) + "\"[^)]*description\\s*=\\s*\")([^\"]+)(\")", Pattern.DOTALL);
-            Matcher funcMatcher = funcPattern.matcher(content);
-            if (funcMatcher.find()) {
-                content = funcMatcher.replaceFirst("$1" + description.replace("\\", "\\\\").replace("$", "\\\\$") + "$3");
-            }
-        }
-
-        // 5. Update register flag
-        Pattern regPattern = Pattern.compile("(@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(toolName) + "\"[^)]*register\\s*=\\s*)(true|false)([^a-zA-Z0-9])", Pattern.DOTALL);
-        Matcher regMatcher = regPattern.matcher(content);
-        if (regMatcher.find()) {
-            content = regMatcher.replaceFirst("$1" + register + "$3");
-        } else {
-            Pattern addRegPattern = Pattern.compile("(@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(toolName) + "\")", Pattern.DOTALL);
-            Matcher addRegMatcher = addRegPattern.matcher(content);
-            if (addRegMatcher.find()) {
-                content = addRegMatcher.replaceFirst("$1, register = " + register);
-            }
-        }
-
-        // 5.5 Update requiresApproval flag
-        if (requiresApproval != null) {
-            Pattern appPattern = Pattern.compile("(@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(toolName) + "\"[^)]*requiresApproval\\s*=\\s*)(true|false)([^a-zA-Z0-9])", Pattern.DOTALL);
-            Matcher appMatcher = appPattern.matcher(content);
-            if (appMatcher.find()) {
-                content = appMatcher.replaceFirst("$1" + requiresApproval + "$3");
-            } else {
-                Pattern addAppPattern = Pattern.compile("(@McpFunction\\s*\\([^)]*name\\s*=\\s*\"" + Pattern.quote(toolName) + "\")", Pattern.DOTALL);
-                Matcher addAppMatcher = addAppPattern.matcher(content);
-                if (addAppMatcher.find()) {
-                    content = addAppMatcher.replaceFirst("$1, requiresApproval = " + requiresApproval);
-                }
-            }
-        }
-
-        // 6. Write back to file
+        String content = Files.readString(targetFile);
+        content = updateMcpTool(content, toolName, description);
+        content = updateToolHint(content, domainGroup, register, requiresApproval);
        Files.writeString(targetFile, content);
    }
+
+    private static Path findToolSource(Path rootDirectory, String toolName) throws IOException {
+        try (var paths = Files.walk(rootDirectory)) {
+            return paths.filter(Files::isRegularFile)
+                    .filter(path -> path.getFileName().toString().endsWith("UseCase.java"))
+                    .filter(path -> path.toString().contains("dap-was-") || path.toString().contains("axhub-tool-"))
+                    .sorted(Comparator.naturalOrder())
+                    .filter(path -> containsMcpTool(path, toolName))
+                    .findFirst()
+                    .orElse(null);
+        }
+    }
+
+    private static boolean containsMcpTool(Path path, String toolName) {
+        try {
+            return annotationArguments(Files.readString(path), "McpTool", toolName) != null;
+        } catch (IOException exception) {
+            throw new IllegalStateException("Failed to read tool source: " + path, exception);
+        }
+    }
+
+    private static String updateMcpTool(String content, String toolName, String description) {
+        AnnotationRange range = annotationArguments(content, "McpTool", toolName);
+        if (range == null) {
+            throw new IllegalArgumentException("McpTool declaration not found: " + toolName);
+        }
+        return description == null ? content : replaceAttribute(content, range, "description", quote(description));
+    }
+
+    private static String updateToolHint(String content, String domainGroup, boolean register, Boolean requiresApproval) {
+        AnnotationRange range = annotationArguments(content, "ToolHint", null);
+        if (range == null) {
+            throw new IllegalArgumentException("ToolHint declaration not found next to McpTool");
+        }
+        String updated = replaceAttribute(content, range, "register", Boolean.toString(register));
+        range = annotationArguments(updated, "ToolHint", null);
+        if (requiresApproval != null) {
+            updated = replaceAttribute(updated, range, "requiresApproval", Boolean.toString(requiresApproval));
+            range = annotationArguments(updated, "ToolHint", null);
+        }
+        if (domainGroup != null && !domainGroup.isBlank()) {
+            updated = replaceAttribute(updated, range, "group", quote(domainGroup));
+        }
+        return updated;
+    }
+
+    private static String replaceAttribute(String content, AnnotationRange range, String attribute, String value) {
+        String arguments = content.substring(range.argumentsStart(), range.argumentsEnd());
+        String pattern = "\\b" + attribute + "\\s*=\\s*(?:true|false|\\\"(?:\\\\.|[^\\\"\\\\])*\\\")";
+        String replacement = arguments.replaceFirst(pattern, attribute + " = " + value);
+        if (replacement.equals(arguments)) {
+            replacement = arguments.isBlank() ? attribute + " = " + value : arguments + ", " + attribute + " = " + value;
+        }
+        return content.substring(0, range.argumentsStart()) + replacement + content.substring(range.argumentsEnd());
+    }
+
+    private static String quote(String value) {
+        return "\"" + value.replace("\\", "\\\\").replace("\"", "\\\"") + "\"";
+    }
+
+    private static AnnotationRange annotationArguments(String content, String annotationName, String toolName) {
+        int offset = content.indexOf("@" + annotationName);
+        while (offset >= 0) {
+            int openingParenthesis = content.indexOf('(', offset);
+            int closingParenthesis = findAnnotationEnd(content, openingParenthesis);
+            if (openingParenthesis < 0 || closingParenthesis < 0) {
+                return null;
+            }
+            AnnotationRange range = new AnnotationRange(openingParenthesis + 1, closingParenthesis);
+            if (toolName == null || content.substring(range.argumentsStart(), range.argumentsEnd())
+                    .matches("(?s).*\\bname\\s*=\\s*\\\"" + java.util.regex.Pattern.quote(toolName) + "\\\".*")) {
+                return range;
+            }
+            offset = content.indexOf("@" + annotationName, closingParenthesis + 1);
+        }
+        return null;
+    }
+
+    private static int findAnnotationEnd(String content, int openingParenthesis) {
+        int depth = 0;
+        boolean inString = false;
+        boolean escaped = false;
+        for (int index = openingParenthesis; index < content.length(); index++) {
+            char character = content.charAt(index);
+            if (inString) {
+                if (escaped) {
+                    escaped = false;
+                } else if (character == '\\') {
+                    escaped = true;
+                } else if (character == '\"') {
+                    inString = false;
+                }
+            } else if (character == '\"') {
+                inString = true;
+            } else if (character == '(') {
+                depth++;
+            } else if (character == ')' && --depth == 0) {
+                return index;
+            }
+        }
+        return -1;
+    }
+
+    private record AnnotationRange(int argumentsStart, int argumentsEnd) {
+    }
 }
--- a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/validation/McpToolNameValidator.java
+++ b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/validation/McpToolNameValidator.java
@@ -15,7 +15,7 @@ import java.util.regex.Pattern;
 /**
 * @package io.shinhanlife.dap.lib.validation
 * @className McpToolNameValidator
- * @description Validates unique MCP function names across tool modules
+ * @description Validates unique MCP SDK tool names across tool modules
 * @author 0986406
 * @create 2026.07.27
 * <pre>
@@ -88,7 +88,7 @@ public final class McpToolNameValidator {
            throw new UncheckedIOException("Failed to read " + source, exception);
        }

-        int annotationOffset = content.indexOf("@McpFunction");
+        int annotationOffset = content.indexOf("@McpTool");
        while (annotationOffset >= 0) {
            int openingParenthesis = content.indexOf('(', annotationOffset);
            int closingParenthesis = findAnnotationEnd(content, openingParenthesis);
@@ -103,7 +103,7 @@ public final class McpToolNameValidator {
                declarationsByName.computeIfAbsent(toolName, ignored -> new ArrayList<>())
                        .add(new ToolDeclaration(moduleName, source, line));
            }
-            annotationOffset = content.indexOf("@McpFunction", closingParenthesis + 1);
+            annotationOffset = content.indexOf("@McpTool", closingParenthesis + 1);
        }
    }

--- a/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/util/ToolSchemaResolverTest.java
+++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/util/ToolSchemaResolverTest.java
@@ -8,26 +8,17 @@ import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
 import org.junit.jupiter.api.Test;
+import org.springframework.ai.mcp.annotation.McpTool;

 class ToolSchemaResolverTest {

    private final ToolSchemaResolver resolver = new ToolSchemaResolver(new ObjectMapper());

-    @Test
-    void usesInlineSchemaBeforeAutomaticDtoSchema() throws Exception {
-        Method method = InlineSchemaTool.class.getDeclaredMethod("search", AutomaticRequest.class);
-
-        Map<String, Object> schema = resolver.resolve(method.getAnnotation(McpFunction.class), AutomaticRequest.class);
-
-        assertEquals(false, schema.get("additionalProperties"));
-        assertTrue(!properties(schema).containsKey("differentField"));
-    }
-
    @Test
    void generatesSchemaFromRequestDtoWhenNoExplicitSchemaIsConfigured() throws Exception {
        Method method = AutomaticSchemaTool.class.getDeclaredMethod("search", AutomaticRequest.class);

-        Map<String, Object> schema = resolver.resolve(method.getAnnotation(McpFunction.class), AutomaticRequest.class);
+        Map<String, Object> schema = resolver.resolve(method.getAnnotation(McpTool.class), null, AutomaticRequest.class);

        assertTrue(properties(schema).containsKey("differentField"));
    }
@@ -35,7 +26,7 @@ class ToolSchemaResolverTest {
    @Test
    void resolvesExplicitOutputSchema() throws Exception {
        Method method = OutputSchemaTool.class.getDeclaredMethod("search", AutomaticRequest.class);
-        Map<String, Object> schema = resolver.resolveOutput(method.getAnnotation(McpFunction.class));
+        Map<String, Object> schema = resolver.resolveOutput(method.getAnnotation(McpTool.class));
        assertEquals(false, schema.get("additionalProperties"));
        assertTrue(properties(schema).containsKey("resultCode"));
    }
@@ -45,7 +36,7 @@ class ToolSchemaResolverTest {
        Method method = AutomaticOutputSchemaTool.class.getDeclaredMethod("search", AutomaticRequest.class);

        Map<String, Object> schema = resolver.resolveOutput(
-                method.getAnnotation(McpFunction.class), SimpleResponse.class);
+                method.getAnnotation(McpTool.class), SimpleResponse.class);

        assertEquals(List.of("resultCode"), schema.get("required"));
        assertEquals(List.of("SUCCESS", "FAILURE"), property(schema, "resultCode").get("enum"));
@@ -55,7 +46,7 @@ class ToolSchemaResolverTest {
    void doesNotEnableOutputValidationWhenOutputSchemaIsNotDeclared() throws Exception {
        Method method = AutomaticSchemaTool.class.getDeclaredMethod("search", AutomaticRequest.class);
        Map<String, Object> schema = resolver.resolveOutput(
-                method.getAnnotation(McpFunction.class), AutomaticRequest.class);
+                method.getAnnotation(McpTool.class), AutomaticRequest.class);
        assertTrue(schema.isEmpty());
    }
    @SuppressWarnings("unchecked")
@@ -68,17 +59,14 @@ class ToolSchemaResolverTest {
        return (Map<String, Object>) properties(schema).get(name);
    }

-    static class InlineSchemaTool {
-        void search(AutomaticRequest request) {
-        }
-    }
-
    static class AutomaticSchemaTool {
+        @McpTool(name = "oth.test.automatic.search")
        void search(AutomaticRequest request) {
        }
    }

    static class AutomaticOutputSchemaTool {
+        @McpTool(name = "oth.test.output.search")
        SimpleResponse search(AutomaticRequest request) {
            return null;
        }
@@ -91,6 +79,7 @@ class ToolSchemaResolverTest {
    }

    static class OutputSchemaTool {
+        @McpTool(name = "oth.test.explicit.search")
        void search(AutomaticRequest request) {
        }
    }
--- a/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/util/ToolSourceUpdaterTest.java
+++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/util/ToolSourceUpdaterTest.java
@@ -0,0 +1,37 @@
+package io.shinhanlife.dap.lib.util;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+class ToolSourceUpdaterTest {
+
+    @TempDir
+    Path temporaryRoot;
+
+    @Test
+    void updatesSdkAndProjectOwnedMetadataOnTheSameToolMethod() throws Exception {
+        Path source = temporaryRoot.resolve("dap-was-oth/src/main/java/example/SampleUseCase.java");
+        Files.createDirectories(source.getParent());
+        Files.writeString(source, """
+                package example;
+                import org.springframework.ai.mcp.annotation.McpTool;
+                import io.shinhanlife.dap.lib.annotation.ToolHint;
+                interface SampleUseCase {
+                    @McpTool(name = "oth.cmm.sample.search", description = "old")
+                    @ToolHint(register = false, requiresApproval = false)
+                    void search();
+                }
+                """);
+
+        ToolSourceUpdater.updateToolSource(temporaryRoot, "oth.cmm.sample.search", "customer", "new", true, true);
+
+        String updated = Files.readString(source);
+        assertTrue(updated.contains("@McpTool(name = \"oth.cmm.sample.search\", description = \"new\")"));
+        assertTrue(updated.contains("@ToolHint(register = true, requiresApproval = true"));
+        assertTrue(updated.contains("group = \"customer\""), updated);
+    }
+}
--- a/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/validation/McpToolNameValidatorTest.java
+++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/validation/McpToolNameValidatorTest.java
@@ -30,7 +30,7 @@ class McpToolNameValidatorTest {
    Path temporaryRoot;

    @Test
-    void rejectsDuplicateMcpFunctionNamesAcrossToolModules() throws IOException {
+    void rejectsDuplicateMcpToolNamesAcrossToolModules() throws IOException {
        writeToolSource("dap-was-first", "FirstTool.java", "first", "oth.sms.notification.send");
        writeToolSource("dap-was-second", "SecondTool.java", "second", "oth.sms.notification.send");

@@ -43,7 +43,7 @@ class McpToolNameValidatorTest {
    }

    @Test
-    void validationRunnerRejectsDuplicateMcpFunctionNamesBeforePackaging() throws IOException {
+    void validationRunnerRejectsDuplicateMcpToolNamesBeforePackaging() throws IOException {
        writeToolSource("dap-was-first", "FirstTool.java", "first", "oth.sms.notification.send");
        writeToolSource("dap-was-second", "SecondTool.java", "second", "oth.sms.notification.send");

@@ -90,10 +90,13 @@ class McpToolNameValidatorTest {
        Files.writeString(source, """
                package example;

+                import org.springframework.ai.mcp.annotation.McpTool;
+
                class %s {
+                    @McpTool(name = "%s")
                    void execute() { }
                }
-                """.formatted(className, className, toolName));
+                """.formatted(className, toolName));
    }

    private Path findProjectRoot() {
--- a/dap-was-lib/src/test/java/io/shinhanlife/dap/mcc/manifest/ToolManifestServiceTest.java
+++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/mcc/manifest/ToolManifestServiceTest.java
@@ -1,4 +1,4 @@
-package io.shinhanlife.dap.mcc.manifest;
+package io.shinhanlife.dap.lib.manifest;

 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
--- a/dap-was-lib/src/test/java/io/shinhanlife/dap/mcc/mcp/ToolMcpServerConfigurationTest.java
+++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/mcc/mcp/ToolMcpServerConfigurationTest.java
@@ -2,6 +2,7 @@ package io.shinhanlife.dap.mcc.mcp;

 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.assertj.core.api.Assertions.assertThat;

 import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTransportProvider;
 import io.shinhanlife.dap.lib.mcp.ToolMcpServerConfiguration;