kimhyungsik/ax_hub_mcp_tool
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: cache tool methods on pod startup
All checks were successful
Deploy to OCIWP / deploy (push) Successful in 1m57s
Details

Browse Source
This commit is contained in:
jade
2026-08-10 09:06:56 +09:00
parent 813c846f05
commit 6809435cc2
4 changed files with 156 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolExecutionService.java
View File

@@ -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<String, Object> arguments, String requestId) {
private ToolExecutionResult validateInput(McpToolMethodRegistry.RegisteredTool tool,
Map<String, Object> arguments, String requestId) {
if (tool.method().getParameterCount() == 0 || Map.class.isAssignableFrom(tool.method().getParameterTypes()[0])) return null;
try {
Map<String, Object> 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<String, Object> 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) { }
}
}

87
dap-was-lib/src/main/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistry.java Normal file
View File

@@ -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<String, RegisteredTool> tools = Map.of();
@EventListener(ApplicationReadyEvent.class)
public void initialize() {
Map<String, RegisteredTool> 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<String, RegisteredTool> discovered, String toolName,
RegisteredTool tool) {
String namespace = mcpProperties.getNamespace();
if (StringUtils.hasText(namespace)) {
register(discovered, namespace + "_" + toolName, tool);
}
}
private void register(Map<String, RegisteredTool> 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) {
}
}

57
dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/McpToolMethodRegistryTest.java Normal file
View File

@@ -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;
}
}
}

4
dap-was-lib/src/test/java/io/shinhanlife/dap/lib/mcp/ToolExecutionServiceTest.java
View File

@@ -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));
}