diff --git a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolExecutionService.java b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolExecutionService.java index bb395f5b..6b0c6cfb 100644 --- a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolExecutionService.java +++ b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolExecutionService.java @@ -2,8 +2,6 @@ package io.shinhanlife.dap.lib.mcp; import com.fasterxml.jackson.databind.ObjectMapper; import com.networknt.schema.Error; -import io.shinhanlife.dap.lib.annotation.ToolHint; -import io.shinhanlife.dap.lib.config.McpProperties; import io.shinhanlife.dap.lib.util.ToolSchemaResolver; import io.shinhanlife.dap.lib.validation.ToolArgumentSchemaValidator; import java.lang.reflect.Method; @@ -12,21 +10,16 @@ import java.util.List; import java.util.Map; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.springaicommunity.mcp.annotation.McpTool; -import org.springframework.aop.support.AopUtils; -import org.springframework.context.ApplicationContext; -import org.springframework.core.annotation.AnnotationUtils; import org.springframework.stereotype.Service; -/** Executes a discovered Tool independently from its HTTP or MCP transport. */ +/** Executes a cached Tool independently from its HTTP or MCP transport. */ @Slf4j @Service @RequiredArgsConstructor public class McpToolExecutionService { - private final ApplicationContext applicationContext; + private final McpToolMethodRegistry toolMethodRegistry; private final ObjectMapper objectMapper; - private final McpProperties mcpProperties; private final ToolArgumentSchemaValidator toolArgumentSchemaValidator; private final ToolSchemaResolver toolSchemaResolver; @@ -37,7 +30,7 @@ public class McpToolExecutionService { String requestId = requestHeaders == null ? null : requestHeaders.requestId(); log.info("[Tool] IN - trace-id: {}, request-id: {}, tool: {}", traceId, requestId, functionName); - ResolvedTool resolvedTool = findTool(functionName); + McpToolMethodRegistry.RegisteredTool resolvedTool = toolMethodRegistry.find(functionName); if (resolvedTool == null) { return error(404, "TOOL_NOT_FOUND", "Tool not found: " + functionName, headerRequestId); } @@ -62,35 +55,8 @@ public class McpToolExecutionService { } } - private ResolvedTool findTool(String functionName) { - for (Object bean : applicationContext.getBeansOfType(Object.class).values()) { - Class targetClass = AopUtils.getTargetClass(bean); - for (Method declaredMethod : targetClass.getDeclaredMethods()) { - McpTool annotation = AnnotationUtils.findAnnotation(declaredMethod, McpTool.class); - if (annotation != null && matches(functionName, annotation.name())) { - return new ResolvedTool(bean, findInvocableMethod(bean, declaredMethod), annotation, - AnnotationUtils.findAnnotation(declaredMethod, ToolHint.class)); - } - } - } - return null; - } - - private boolean matches(String requestedName, String baseName) { - String namespace = mcpProperties.getNamespace(); - String expectedName = namespace != null && !namespace.isEmpty() ? namespace + "_" + baseName : baseName; - return expectedName.equals(requestedName) || baseName.equals(requestedName); - } - - private Method findInvocableMethod(Object bean, Method declaredMethod) { - try { - return bean.getClass().getMethod(declaredMethod.getName(), declaredMethod.getParameterTypes()); - } catch (NoSuchMethodException ignored) { - return declaredMethod; - } - } - - private ToolExecutionResult validateInput(ResolvedTool tool, Map arguments, String requestId) { + private ToolExecutionResult validateInput(McpToolMethodRegistry.RegisteredTool tool, + Map arguments, String requestId) { if (tool.method().getParameterCount() == 0 || Map.class.isAssignableFrom(tool.method().getParameterTypes()[0])) return null; try { Map schema = toolSchemaResolver.resolve(tool.annotation(), tool.hint(), tool.method().getParameterTypes()[0]); @@ -107,11 +73,12 @@ public class McpToolExecutionService { return objectMapper.convertValue(arguments, method.getParameterTypes()[0]); } - private Object invoke(ResolvedTool tool, Object argument) throws Exception { + private Object invoke(McpToolMethodRegistry.RegisteredTool tool, Object argument) throws Exception { return tool.method().getParameterCount() == 0 ? tool.method().invoke(tool.bean()) : tool.method().invoke(tool.bean(), argument); } - private ToolExecutionResult validateOutput(ResolvedTool tool, Object methodResult, String requestId) { + private ToolExecutionResult validateOutput(McpToolMethodRegistry.RegisteredTool tool, + Object methodResult, String requestId) { try { Map outputSchema = toolSchemaResolver.resolveOutput(tool.annotation(), tool.method().getReturnType(), tool.hint()); if (!outputSchema.isEmpty() && !toolArgumentSchemaValidator.validateValue(outputSchema, methodResult).isEmpty()) { @@ -131,6 +98,4 @@ public class McpToolExecutionService { if (requestId != null) body.put("request_id", requestId); return new ToolExecutionResult(statusCode, body, Map.of()); } - - private record ResolvedTool(Object bean, Method method, McpTool annotation, ToolHint hint) { } -} +} \ No newline at end of file diff --git a/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistry.java b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistry.java new file mode 100644 index 00000000..cd3dc4a1 --- /dev/null +++ b/dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistry.java @@ -0,0 +1,87 @@ +package io.shinhanlife.dap.lib.mcp; + +import io.shinhanlife.dap.lib.annotation.ToolHint; +import io.shinhanlife.dap.lib.config.McpProperties; +import java.lang.reflect.Method; +import java.util.LinkedHashMap; +import java.util.Map; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springaicommunity.mcp.annotation.McpTool; +import org.springframework.aop.support.AopUtils; +import org.springframework.boot.context.event.ApplicationReadyEvent; +import org.springframework.context.ApplicationContext; +import org.springframework.context.event.EventListener; +import org.springframework.core.annotation.AnnotationUtils; +import org.springframework.stereotype.Component; +import org.springframework.util.StringUtils; + +/** + * Caches executable {@link McpTool} methods once when a Tool Pod starts. + */ +@Slf4j +@Component +@RequiredArgsConstructor +public class McpToolMethodRegistry { + + private final ApplicationContext applicationContext; + private final McpProperties mcpProperties; + + private volatile Map tools = Map.of(); + + @EventListener(ApplicationReadyEvent.class) + public void initialize() { + Map discovered = new LinkedHashMap<>(); + + for (Object bean : applicationContext.getBeansOfType(Object.class).values()) { + Class targetClass = AopUtils.getTargetClass(bean); + for (Method declaredMethod : targetClass.getDeclaredMethods()) { + McpTool annotation = AnnotationUtils.findAnnotation(declaredMethod, McpTool.class); + if (annotation == null) { + continue; + } + + RegisteredTool tool = new RegisteredTool( + bean, + findInvocableMethod(bean, declaredMethod), + annotation, + AnnotationUtils.findAnnotation(declaredMethod, ToolHint.class)); + register(discovered, annotation.name(), tool); + registerNamespaceAlias(discovered, annotation.name(), tool); + } + } + + tools = Map.copyOf(discovered); + log.info("[Tool Registry] {} executable tool names cached", tools.size()); + } + + public RegisteredTool find(String toolName) { + return tools.get(toolName); + } + + private void registerNamespaceAlias(Map discovered, String toolName, + RegisteredTool tool) { + String namespace = mcpProperties.getNamespace(); + if (StringUtils.hasText(namespace)) { + register(discovered, namespace + "_" + toolName, tool); + } + } + + private void register(Map discovered, String toolName, RegisteredTool tool) { + RegisteredTool existing = discovered.putIfAbsent(toolName, tool); + if (existing != null && existing != tool) { + throw new IllegalStateException("Duplicate @McpTool name: " + toolName); + } + } + + private Method findInvocableMethod(Object bean, Method declaredMethod) { + try { + return bean.getClass().getMethod(declaredMethod.getName(), declaredMethod.getParameterTypes()); + } catch (NoSuchMethodException ignored) { + return declaredMethod; + } + } + + public record RegisteredTool(Object bean, Method method, McpTool annotation, ToolHint hint) { + } +} \ No newline at end of file diff --git a/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistryTest.java b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistryTest.java new file mode 100644 index 00000000..050d6b10 --- /dev/null +++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistryTest.java @@ -0,0 +1,57 @@ +package io.shinhanlife.dap.lib.mcp; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import io.shinhanlife.dap.lib.config.McpProperties; +import org.junit.jupiter.api.Test; +import org.springaicommunity.mcp.annotation.McpTool; +import org.springframework.context.support.StaticApplicationContext; + +class McpToolMethodRegistryTest { + + @Test + void indexesAnnotatedToolDuringInitialization() { + StaticApplicationContext context = contextWith(new EchoTool()); + McpToolMethodRegistry registry = new McpToolMethodRegistry(context, new McpProperties()); + + registry.initialize(); + + McpToolMethodRegistry.RegisteredTool tool = registry.find("oth.cmm.echo.search"); + assertNotNull(tool); + assertEquals("execute", tool.method().getName()); + assertEquals("oth.cmm.echo.search", tool.annotation().name()); + } + + @Test + void rejectsDuplicateToolNamesDuringInitialization() { + StaticApplicationContext context = contextWith(new EchoTool(), new DuplicateEchoTool()); + McpToolMethodRegistry registry = new McpToolMethodRegistry(context, new McpProperties()); + + assertThrows(IllegalStateException.class, registry::initialize); + } + + private StaticApplicationContext contextWith(Object... tools) { + StaticApplicationContext context = new StaticApplicationContext(); + for (int index = 0; index < tools.length; index++) { + context.getBeanFactory().registerSingleton("tool" + index, tools[index]); + } + context.refresh(); + return context; + } + + static class EchoTool { + @McpTool(name = "oth.cmm.echo.search") + public String execute(String request) { + return request; + } + } + + static class DuplicateEchoTool { + @McpTool(name = "oth.cmm.echo.search") + public String execute(String request) { + return request; + } + } +} \ No newline at end of file diff --git a/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/ToolExecutionServiceTest.java b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/ToolExecutionServiceTest.java index 91f98e32..26ddc40a 100644 --- a/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/ToolExecutionServiceTest.java +++ b/dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/ToolExecutionServiceTest.java @@ -39,7 +39,9 @@ class McpToolExecutionServiceTest { StaticApplicationContext context = new StaticApplicationContext(); context.getBeanFactory().registerSingleton("toolBean", toolBean); context.refresh(); - return new McpToolExecutionService(context, objectMapper, new McpProperties(), + McpToolMethodRegistry registry = new McpToolMethodRegistry(context, new McpProperties()); + registry.initialize(); + return new McpToolExecutionService(registry, objectMapper, new ToolArgumentSchemaValidator(objectMapper), new ToolSchemaResolver(objectMapper)); }