From 78abea702f9bc17b30b0d830dc23591086c040ce Mon Sep 17 00:00:00 2001 From: hengyunabc Date: Fri, 26 Dec 2025 20:29:56 +0800 Subject: [PATCH] arthas-mcp-integration-test (#3115) --- .github/workflows/test.yaml | 4 +- arthas-mcp-integration-test/README.md | 21 + arthas-mcp-integration-test/pom.xml | 74 ++ .../arthas/mcp/it/ArthasMcpJavaSdkIT.java | 654 ++++++++++++++++++ .../arthas/mcp/it/ArthasMcpToolsIT.java | 401 +++++++++++ .../taobao/arthas/mcp/it/TargetJvmApp.java | 25 + .../mcp/tool/function/AbstractArthasTool.java | 79 ++- .../tool/function/StreamableToolUtils.java | 10 +- pom.xml | 16 + 9 files changed, 1277 insertions(+), 7 deletions(-) create mode 100644 arthas-mcp-integration-test/README.md create mode 100644 arthas-mcp-integration-test/pom.xml create mode 100644 arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpJavaSdkIT.java create mode 100644 arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpToolsIT.java create mode 100644 arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/TargetJvmApp.java diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index e01b3c6e5..8d50ef44f 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -7,7 +7,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - java: [8, 11, 17, 19] + java: [8, 11, 17, 21, 25] steps: - uses: actions/checkout@v3 - name: Setup java @@ -17,7 +17,7 @@ jobs: distribution: "zulu" cache: "maven" - name: Build with Maven - run: mvn -V -ntp clean install -P full + run: mvn -V -ntp clean install -P full verify windows_build: runs-on: windows-2022 diff --git a/arthas-mcp-integration-test/README.md b/arthas-mcp-integration-test/README.md new file mode 100644 index 000000000..c71c81ac9 --- /dev/null +++ b/arthas-mcp-integration-test/README.md @@ -0,0 +1,21 @@ +# arthas-mcp-integration-test + +本模块提供 Arthas MCP Server 的集成测试： + +- 测试会启动一个独立的目标 JVM（`TargetJvmApp`）。 +- 通过 `packaging/target/arthas-bin/as.sh` 动态 attach 到目标 JVM，在目标 JVM 内启动 Arthas Server（仅开启 HTTP 端口，telnet 端口设置为 0）。 +- 使用最小 MCP（Streamable HTTP + SSE）客户端调用 `tools/list` 与 `tools/call`，验证 MCP tools 功能可用。 + +## 运行方式 + +在项目根目录执行： + +```bash +./mvnw -pl arthas-mcp-integration-test -am verify +``` + +说明： + +- `-am` 会确保 `packaging` 等依赖模块先构建，从而在 `packaging/target/arthas-bin` 生成可用的 `as.sh` 与相关 jar。 +- 该集成测试依赖本机 `bash`，Windows 环境会自动跳过。 + diff --git a/arthas-mcp-integration-test/pom.xml b/arthas-mcp-integration-test/pom.xml new file mode 100644 index 000000000..ee055d819 --- /dev/null +++ b/arthas-mcp-integration-test/pom.xml @@ -0,0 +1,74 @@ + + + 4.0.0 + + + com.taobao.arthas + arthas-all + ${revision} + ../pom.xml + + + arthas-mcp-integration-test + arthas-mcp-integration-test + https://github.com/alibaba/arthas + + + + com.taobao.arthas + arthas-packaging + ${project.version} + test + + + + com.taobao.arthas + arthas-mcp-server + ${project.version} + test + + + + org.junit.jupiter + junit-jupiter + test + + + org.assertj + assertj-core + test + + + + io.modelcontextprotocol.sdk + mcp + + + + + + + org.apache.maven.plugins + maven-surefire-plugin + 3.2.5 + + + it-verify + verify + + test + + + + **/*IT.java + + + + + + + + + diff --git a/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpJavaSdkIT.java b/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpJavaSdkIT.java new file mode 100644 index 000000000..96dddbe5d --- /dev/null +++ b/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpJavaSdkIT.java @@ -0,0 +1,654 @@ +package com.taobao.arthas.mcp.it; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.taobao.arthas.mcp.server.util.JsonParser; +import io.modelcontextprotocol.client.McpClient; +import io.modelcontextprotocol.client.McpSyncClient; +import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; +import io.modelcontextprotocol.spec.McpClientTransport; +import io.modelcontextprotocol.spec.McpSchema; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.net.InetSocketAddress; +import java.net.Socket; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@Execution(ExecutionMode.SAME_THREAD) +class ArthasMcpJavaSdkIT { + + private static final ObjectMapper OBJECT_MAPPER = JsonParser.getObjectMapper(); + + private static final String TARGET_CLASS_PATTERN = TargetJvmApp.class.getName(); + + private static final String TARGET_METHOD_PATTERN = "hotMethod"; + + private static final List EXPECTED_TOOL_NAMES = Arrays.asList( + "classloader", + "dashboard", + "dump", + "getstatic", + "heapdump", + "jad", + "jvm", + "mc", + "mbean", + "memory", + "monitor", + "ognl", + "options", + "perfcounter", + "redefine", + "retransform", + "sc", + "sm", + "stack", + "stop", + "sysenv", + "sysprop", + "thread", + "trace", + "tt", + "vmoption", + "vmtool", + "watch" + ); + + private Environment env; + + @BeforeAll + void setUp() throws Exception { + Assumptions.assumeFalse(isWindows(), "集成测试依赖 bash/as.sh，Windows 环境跳过"); + this.env = Environment.start("arthas-mcp-java-sdk-it", "arthas-mcp-java-sdk-it-home"); + } + + @AfterAll + void tearDown() { + if (this.env != null) { + this.env.close(); + this.env = null; + } + } + + @Test + @Timeout(value = 3, unit = TimeUnit.MINUTES) + void should_list_all_mcp_tools_via_java_mcp_sdk() { + Set expected = new HashSet(EXPECTED_TOOL_NAMES); + Set actual = new HashSet(this.env.toolNames); + + Set missing = new HashSet(expected); + missing.removeAll(actual); + + Set extra = new HashSet(actual); + extra.removeAll(expected); + + assertThat(missing).as("tools/list 缺少工具: %s", missing).isEmpty(); + assertThat(extra).as("tools/list 存在未覆盖的工具: %s", extra).isEmpty(); + } + + static Stream toolNamesExceptStop() { + List toolNames = new ArrayList(EXPECTED_TOOL_NAMES); + toolNames.remove("stop"); + return toolNames.stream(); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("toolNamesExceptStop") + @Timeout(value = 3, unit = TimeUnit.MINUTES) + void should_call_each_mcp_tool_via_java_mcp_sdk(String toolName) throws Exception { + Map args = createArgumentsForTool(toolName, this.env); + McpSchema.CallToolResult result = this.env.client.callTool(new McpSchema.CallToolRequest(toolName, args)); + + String body = assertCallToolSuccess(toolName, result); + assertToolSideEffects(toolName, this.env, body); + } + + @Test + @Timeout(value = 3, unit = TimeUnit.MINUTES) + void should_call_stop_tool_via_java_mcp_sdk() throws Exception { + Assumptions.assumeFalse(isWindows(), "集成测试依赖 bash/as.sh，Windows 环境跳过"); + + Environment stopEnv = null; + try { + stopEnv = Environment.start("arthas-mcp-java-sdk-stop-it", "arthas-mcp-java-sdk-stop-it-home"); + Map args = new HashMap(); + args.put("delayMs", 200); + + McpSchema.CallToolResult result = stopEnv.client.callTool(new McpSchema.CallToolRequest("stop", args)); + assertCallToolSuccess("stop", result); + + waitForPortClosed("127.0.0.1", stopEnv.httpPort, Duration.ofSeconds(15)); + } finally { + if (stopEnv != null) { + stopEnv.close(); + } + } + } + + private static void assertToolSideEffects(String toolName, Environment env, String body) throws Exception { + if ("heapdump".equals(toolName)) { + Path heapdumpFile = env.tempHome.resolve("heapdump.hprof"); + assertThat(heapdumpFile).exists(); + assertThat(Files.size(heapdumpFile)).isGreaterThan(0L); + try { + Files.deleteIfExists(heapdumpFile); + } catch (Exception ignored) { + } + return; + } + + if ("dump".equals(toolName)) { + Path dumpOutputDir = env.tempHome.resolve("dump-output"); + assertThat(dumpOutputDir).isDirectory(); + assertThat(countFilesWithSuffix(dumpOutputDir, ".class")).isGreaterThan(0); + return; + } + + if ("mc".equals(toolName)) { + Path mcOutputDir = env.tempHome.resolve("mc-output"); + assertThat(mcOutputDir).isDirectory(); + assertThat(countFilesWithSuffix(mcOutputDir, ".class")).isGreaterThan(0); + return; + } + + if (isStreamableTool(toolName)) { + JsonNode node = OBJECT_MAPPER.readTree(body); + if (node != null && node.isObject()) { + JsonNode resultCount = node.get("resultCount"); + if (resultCount != null && resultCount.canConvertToInt()) { + int count = resultCount.asInt(); + assertThat(count).as("tool=%s resultCount, body=%s", toolName, body).isGreaterThan(0); + } else { + Assertions.fail("streamable tool 未返回 resultCount: tool=" + toolName + ", body=" + body); + } + } + } + } + + private static boolean isStreamableTool(String toolName) { + return "dashboard".equals(toolName) + || "monitor".equals(toolName) + || "watch".equals(toolName) + || "trace".equals(toolName) + || "stack".equals(toolName) + || "tt".equals(toolName); + } + + private static int countFilesWithSuffix(Path dir, String suffix) throws IOException { + if (dir == null || !Files.isDirectory(dir)) { + return 0; + } + try (Stream stream = Files.walk(dir)) { + return (int) stream + .filter(p -> Files.isRegularFile(p) + && p.getFileName() != null + && p.getFileName().toString().endsWith(suffix)) + .count(); + } + } + + private static Map createArgumentsForTool(String toolName, Environment env) throws IOException { + Map args = new HashMap(); + + if ("jvm".equals(toolName) + || "thread".equals(toolName) + || "memory".equals(toolName) + || "options".equals(toolName) + || "vmoption".equals(toolName) + || "classloader".equals(toolName) + || "perfcounter".equals(toolName)) { + return args; + } + + if ("jad".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + return args; + } + + if ("sc".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + return args; + } + + if ("sm".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + args.put("methodPattern", TARGET_METHOD_PATTERN); + return args; + } + + if ("dump".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + Path outDir = env.tempHome.resolve("dump-output"); + Files.createDirectories(outDir); + args.put("outputDir", outDir.toString()); + args.put("limit", 1); + return args; + } + + if ("mc".equals(toolName)) { + Path sourceFile = env.ensureMcSourceFile(); + Path outDir = env.tempHome.resolve("mc-output"); + Files.createDirectories(outDir); + args.put("javaFilePaths", sourceFile.toString()); + args.put("outputDir", outDir.toString()); + return args; + } + + if ("retransform".equals(toolName) || "redefine".equals(toolName)) { + args.put("classFilePaths", env.targetClassFile.toString()); + return args; + } + + if ("getstatic".equals(toolName)) { + args.put("className", "java.lang.Integer"); + args.put("fieldName", "MAX_VALUE"); + return args; + } + + if ("ognl".equals(toolName)) { + args.put("expression", "@java.lang.System@getProperty(\"java.version\")"); + args.put("expandLevel", 1); + return args; + } + + if ("mbean".equals(toolName)) { + args.put("namePattern", "java.lang:type=Runtime"); + args.put("attributePattern", "Uptime"); + return args; + } + + if ("sysenv".equals(toolName)) { + args.put("envName", "PATH"); + return args; + } + + if ("sysprop".equals(toolName)) { + args.put("propertyName", "java.version"); + return args; + } + + if ("vmtool".equals(toolName)) { + args.put("action", "getInstances"); + args.put("className", TARGET_CLASS_PATTERN); + args.put("limit", 1); + args.put("expandLevel", 1); + args.put("express", "instances.length"); + return args; + } + + if ("heapdump".equals(toolName)) { + Path heapdumpFile = env.tempHome.resolve("heapdump.hprof"); + args.put("filePath", heapdumpFile.toString()); + return args; + } + + if ("dashboard".equals(toolName)) { + args.put("intervalMs", 200); + args.put("numberOfExecutions", 1); + return args; + } + + if ("watch".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + args.put("methodPattern", TARGET_METHOD_PATTERN); + args.put("numberOfExecutions", 1); + args.put("timeout", 10); + return args; + } + + if ("trace".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + args.put("methodPattern", TARGET_METHOD_PATTERN); + args.put("numberOfExecutions", 1); + args.put("timeout", 10); + return args; + } + + if ("monitor".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + args.put("methodPattern", TARGET_METHOD_PATTERN); + args.put("intervalMs", 1000); + args.put("numberOfExecutions", 1); + args.put("timeout", 15); + return args; + } + + if ("stack".equals(toolName)) { + args.put("classPattern", TARGET_CLASS_PATTERN); + args.put("methodPattern", TARGET_METHOD_PATTERN); + args.put("numberOfExecutions", 1); + // CI 环境下 stack 增强+触发可能更慢，适当放大超时时间以减少偶发失败 + args.put("timeout", 30); + return args; + } + + if ("tt".equals(toolName)) { + args.put("action", "record"); + args.put("classPattern", TARGET_CLASS_PATTERN); + args.put("methodPattern", TARGET_METHOD_PATTERN); + args.put("numberOfExecutions", 1); + args.put("timeout", 10); + return args; + } + + throw new IllegalArgumentException("未为 tool 配置参数: " + toolName); + } + + private static String assertCallToolSuccess(String toolName, McpSchema.CallToolResult result) throws Exception { + assertThat(result).as("tool=%s", toolName).isNotNull(); + assertThat(result.isError()).as("tool=%s, content=%s", toolName, result.content()).isNotEqualTo(Boolean.TRUE); + assertThat(result.content()).as("tool=%s", toolName).isNotNull().isNotEmpty(); + + String text = extractTextContent(result); + assertThat(text).as("tool=%s", toolName).isNotBlank(); + + JsonNode node = OBJECT_MAPPER.readTree(text); + if (node != null && node.isObject()) { + JsonNode error = node.get("error"); + if (error != null && error.isBoolean() && error.booleanValue()) { + Assertions.fail("tool 执行返回 error=true: tool=" + toolName + ", body=" + text); + } + JsonNode status = node.get("status"); + if (status != null && status.isTextual() && "error".equalsIgnoreCase(status.asText())) { + Assertions.fail("tool 执行返回 status=error: tool=" + toolName + ", body=" + text); + } + } + + return text; + } + + private static String extractTextContent(McpSchema.CallToolResult result) { + StringBuilder sb = new StringBuilder(); + for (McpSchema.Content content : result.content()) { + if (content instanceof McpSchema.TextContent) { + String text = ((McpSchema.TextContent) content).text(); + if (text != null) { + if (sb.length() > 0) { + sb.append('\n'); + } + sb.append(text); + } + } + } + return sb.toString(); + } + + private static final class Environment implements AutoCloseable { + private final Path arthasHome; + private final Path tempHome; + private final int httpPort; + private final Process targetJvm; + private final McpSyncClient client; + private final Set toolNames; + private final Path targetClassFile; + private Path mcSourceFile; + + private Environment(Path arthasHome, Path tempHome, int httpPort, Process targetJvm, + McpSyncClient client, Set toolNames, Path targetClassFile) { + this.arthasHome = arthasHome; + this.tempHome = tempHome; + this.httpPort = httpPort; + this.targetJvm = targetJvm; + this.client = client; + this.toolNames = toolNames; + this.targetClassFile = targetClassFile; + } + + static Environment start(String clientName, String tempDirPrefix) throws Exception { + Path arthasHome = resolveArthasBinDir(); + assertThat(arthasHome).isDirectory(); + assertThat(arthasHome.resolve("as.sh")).exists(); + assertThat(arthasHome.resolve("arthas-core.jar")).exists(); + assertThat(arthasHome.resolve("arthas-agent.jar")).exists(); + + int telnetPort = 0; + int httpPort = findFreePort(); + + Path tempHome = Files.createTempDirectory(tempDirPrefix); + Process targetJvm = null; + McpSyncClient client = null; + try { + Path targetLog = tempHome.resolve("target-jvm.log"); + targetJvm = startTargetJvm(tempHome, targetLog); + long targetPid = ProcessPid.pidOf(targetJvm); + + Path attachLog = tempHome.resolve("attach.log"); + runAttach(arthasHome, tempHome, attachLog, targetPid, telnetPort, httpPort); + + waitForPortOpen("127.0.0.1", httpPort, Duration.ofSeconds(30)); + + McpClientTransport transport = HttpClientStreamableHttpTransport.builder("http://127.0.0.1:" + httpPort).build(); + client = McpClient.sync(transport) + .clientInfo(new McpSchema.Implementation(clientName, "1.0.0")) + .requestTimeout(Duration.ofSeconds(120)) + .initializationTimeout(Duration.ofSeconds(10)) + .build(); + + McpSchema.InitializeResult initResult = client.initialize(); + assertThat(initResult).isNotNull(); + + McpSchema.ListToolsResult toolsResult = client.listTools(); + assertThat(toolsResult).isNotNull(); + assertThat(toolsResult.tools()).isNotNull(); + + Set toolNames = new HashSet(); + for (McpSchema.Tool tool : toolsResult.tools()) { + toolNames.add(tool.name()); + } + + Path targetClassFile = resolveTargetJvmAppClassFile(); + assertThat(targetClassFile).exists(); + + return new Environment(arthasHome, tempHome, httpPort, targetJvm, client, toolNames, targetClassFile); + } catch (Exception e) { + if (client != null) { + try { + client.closeGracefully(); + } catch (Exception ignored) { + } + } + if (targetJvm != null) { + targetJvm.destroy(); + if (!targetJvm.waitFor(5, TimeUnit.SECONDS)) { + targetJvm.destroyForcibly(); + } + } + deleteDirectoryQuietly(tempHome); + throw e; + } + } + + Path ensureMcSourceFile() throws IOException { + if (this.mcSourceFile != null) { + return this.mcSourceFile; + } + Path source = this.tempHome.resolve("McpMcTestClass.java"); + String code = "" + + "public class McpMcTestClass {\n" + + " public static int add(int a, int b) {\n" + + " return a + b;\n" + + " }\n" + + "}\n"; + Files.write(source, code.getBytes(StandardCharsets.UTF_8)); + this.mcSourceFile = source; + return source; + } + + @Override + public void close() { + if (this.client != null) { + try { + this.client.closeGracefully(); + } catch (Exception ignored) { + } + } + if (this.targetJvm != null) { + this.targetJvm.destroy(); + try { + if (!this.targetJvm.waitFor(5, TimeUnit.SECONDS)) { + this.targetJvm.destroyForcibly(); + } + } catch (Exception ignored) { + } + } + deleteDirectoryQuietly(this.tempHome); + } + } + + private static void runAttach(Path arthasHome, Path tempHome, Path attachLog, long targetPid, + int telnetPort, int httpPort) throws Exception { + ProcessBuilder pb = new ProcessBuilder( + "bash", + arthasHome.resolve("as.sh").toString(), + "--attach-only", + "--arthas-home", arthasHome.toString(), + "--target-ip", "127.0.0.1", + "--telnet-port", String.valueOf(telnetPort), + "--http-port", String.valueOf(httpPort), + String.valueOf(targetPid) + ); + pb.redirectErrorStream(true); + pb.redirectOutput(attachLog.toFile()); + pb.environment().put("JAVA_HOME", System.getProperty("java.home")); + pb.environment().put("HOME", tempHome.toAbsolutePath().toString()); + + Process attach = pb.start(); + if (!attach.waitFor(90, TimeUnit.SECONDS)) { + attach.destroyForcibly(); + Assertions.fail("as.sh attach 超时: " + attachLog); + } + if (attach.exitValue() != 0) { + Assertions.fail("as.sh attach 失败(exit=" + attach.exitValue() + "): " + attachLog); + } + } + + private static Process startTargetJvm(Path workDir, Path targetLog) throws IOException { + String javaBin = Paths.get(System.getProperty("java.home"), "bin", "java").toString(); + String classpath = Paths.get(System.getProperty("basedir"), "target", "test-classes").toString(); + ProcessBuilder pb = new ProcessBuilder(javaBin, "-cp", classpath, TargetJvmApp.class.getName()); + if (workDir != null) { + pb.directory(workDir.toFile()); + } + pb.redirectErrorStream(true); + pb.redirectOutput(targetLog.toFile()); + return pb.start(); + } + + private static Path resolveArthasBinDir() { + String basedir = System.getProperty("basedir"); + assertThat(basedir).as("Maven surefire/failsafe should set system property 'basedir'").isNotBlank(); + return Paths.get(basedir).resolve("../packaging/target/arthas-bin").normalize(); + } + + private static Path resolveTargetJvmAppClassFile() { + String basedir = System.getProperty("basedir"); + assertThat(basedir).as("Maven surefire/failsafe should set system property 'basedir'").isNotBlank(); + return Paths.get(basedir, "target", "test-classes", "com", "taobao", "arthas", "mcp", "it", "TargetJvmApp.class").normalize(); + } + + private static int findFreePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + socket.setReuseAddress(true); + return socket.getLocalPort(); + } + } + + private static void waitForPortOpen(String host, int port, Duration timeout) throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (System.nanoTime() < deadline) { + try (Socket socket = new Socket()) { + socket.connect(new InetSocketAddress(host, port), 500); + return; + } catch (IOException ignored) { + Thread.sleep(200); + } + } + throw new IllegalStateException("等待端口监听超时: " + host + ":" + port); + } + + private static void waitForPortClosed(String host, int port, Duration timeout) throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (System.nanoTime() < deadline) { + try (Socket socket = new Socket()) { + socket.connect(new InetSocketAddress(host, port), 500); + Thread.sleep(200); + } catch (IOException ignored) { + return; + } + } + throw new IllegalStateException("等待端口关闭超时: " + host + ":" + port); + } + + private static boolean isWindows() { + String os = System.getProperty("os.name"); + return os != null && os.toLowerCase(Locale.ROOT).contains("win"); + } + + private static void deleteDirectoryQuietly(Path dir) { + try { + if (!Files.exists(dir)) { + return; + } + Files.walk(dir) + .sorted(Comparator.reverseOrder()) + .map(Path::toFile) + .forEach(File::delete); + } catch (Exception ignored) { + } + } + + private static final class ProcessPid { + private ProcessPid() { + } + + static long pidOf(Process process) { + // Java 9+ 获取 pid + try { + return (Long) Process.class.getMethod("pid").invoke(process); + } catch (Exception ignored) { + // Java 8 兼容处理 + } + try { + java.lang.reflect.Field pidField = process.getClass().getDeclaredField("pid"); + pidField.setAccessible(true); + Object value = pidField.get(process); + if (value instanceof Number) { + return ((Number) value).longValue(); + } + throw new IllegalStateException("Unsupported pid field type: " + value); + } catch (Exception e) { + throw new IllegalStateException("无法获取目标 JVM pid", e); + } + } + } +} diff --git a/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpToolsIT.java b/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpToolsIT.java new file mode 100644 index 000000000..9bf420f35 --- /dev/null +++ b/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/ArthasMcpToolsIT.java @@ -0,0 +1,401 @@ +package com.taobao.arthas.mcp.it; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.taobao.arthas.mcp.server.protocol.spec.HttpHeaders; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.io.BufferedReader; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.net.HttpURLConnection; +import java.net.InetSocketAddress; +import java.net.Socket; +import java.net.URL; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.time.Duration; +import java.util.*; +import java.util.concurrent.TimeUnit; + +import static org.assertj.core.api.Assertions.assertThat; + +class ArthasMcpToolsIT { + + private static final ObjectMapper OBJECT_MAPPER = JsonParser.getObjectMapper(); + + @Test + @Timeout(value = 3, unit = TimeUnit.MINUTES) + void should_list_tools_and_call_tool_via_mcp() throws Exception { + Assumptions.assumeFalse(isWindows(), "集成测试依赖 bash/as.sh，Windows 环境跳过"); + + Path arthasHome = resolveArthasBinDir(); + assertThat(arthasHome).isDirectory(); + assertThat(arthasHome.resolve("as.sh")).exists(); + assertThat(arthasHome.resolve("arthas-core.jar")).exists(); + assertThat(arthasHome.resolve("arthas-agent.jar")).exists(); + + int telnetPort = 0; + int httpPort = findFreePort(); + + Process targetJvm = null; + Path tempHome = null; + try { + tempHome = Files.createTempDirectory("arthas-mcp-it-home"); + + Path targetLog = tempHome.resolve("target-jvm.log"); + targetJvm = startTargetJvm(targetLog); + long targetPid = ProcessPid.pidOf(targetJvm); + + Path attachLog = tempHome.resolve("attach.log"); + runAttach(arthasHome, tempHome, attachLog, targetPid, telnetPort, httpPort); + + waitForPortOpen("127.0.0.1", httpPort, Duration.ofSeconds(30)); + + StreamableMcpHttpClient client = new StreamableMcpHttpClient("127.0.0.1", httpPort, "/mcp"); + String sessionId = retry(Duration.ofSeconds(30), client::initialize); + client.sendInitializedNotification(sessionId); + + McpSchema.ListToolsResult toolsResult = client.listTools(sessionId); + assertThat(toolsResult.getTools()).isNotNull(); + assertThat(toolsResult.getTools().size()).isGreaterThanOrEqualTo(10); + + Set toolNames = new HashSet<>(); + for (McpSchema.Tool tool : toolsResult.getTools()) { + toolNames.add(tool.getName()); + } + assertThat(toolNames).contains("jvm", "jad", "thread"); + + McpSchema.CallToolResult callToolResult = client.callTool(sessionId, "jvm", Collections.emptyMap()); + assertThat(callToolResult).isNotNull(); + assertThat(callToolResult.getIsError()).isNotEqualTo(Boolean.TRUE); + assertThat(callToolResult.getContent()).isNotNull().isNotEmpty(); + } finally { + if (targetJvm != null) { + targetJvm.destroy(); + if (!targetJvm.waitFor(5, TimeUnit.SECONDS)) { + targetJvm.destroyForcibly(); + } + } + if (tempHome != null) { + deleteDirectoryQuietly(tempHome); + } + } + } + + private static String retry(Duration timeout, IoSupplier supplier) throws Exception { + long deadline = System.nanoTime() + timeout.toNanos(); + Exception last = null; + while (System.nanoTime() < deadline) { + try { + return supplier.get(); + } catch (Exception e) { + last = e; + Thread.sleep(200); + } + } + if (last != null) { + throw last; + } + throw new IllegalStateException("重试超时"); + } + + @FunctionalInterface + private interface IoSupplier { + T get() throws Exception; + } + + private static void runAttach(Path arthasHome, Path tempHome, Path attachLog, long targetPid, + int telnetPort, int httpPort) throws Exception { + List command = new ArrayList<>(); + command.add("bash"); + command.add(arthasHome.resolve("as.sh").toString()); + command.add("--attach-only"); + command.add("--arthas-home"); + command.add(arthasHome.toString()); + command.add("--target-ip"); + command.add("127.0.0.1"); + command.add("--telnet-port"); + command.add(String.valueOf(telnetPort)); + command.add("--http-port"); + command.add(String.valueOf(httpPort)); + command.add(String.valueOf(targetPid)); + + ProcessBuilder pb = new ProcessBuilder(command); + pb.redirectErrorStream(true); + pb.redirectOutput(attachLog.toFile()); + Map env = pb.environment(); + env.put("JAVA_HOME", System.getProperty("java.home")); + env.put("HOME", tempHome.toAbsolutePath().toString()); + + Process attach = pb.start(); + if (!attach.waitFor(90, TimeUnit.SECONDS)) { + attach.destroyForcibly(); + Assertions.fail("as.sh attach 超时: " + attachLog); + } + if (attach.exitValue() != 0) { + Assertions.fail("as.sh attach 失败(exit=" + attach.exitValue() + "): " + attachLog); + } + } + + private static Process startTargetJvm(Path targetLog) throws IOException { + String javaBin = Paths.get(System.getProperty("java.home"), "bin", "java").toString(); + String classpath = Paths.get(System.getProperty("basedir"), "target", "test-classes").toString(); + ProcessBuilder pb = new ProcessBuilder(javaBin, "-cp", classpath, TargetJvmApp.class.getName()); + pb.redirectErrorStream(true); + pb.redirectOutput(targetLog.toFile()); + return pb.start(); + } + + private static Path resolveArthasBinDir() { + String basedir = System.getProperty("basedir"); + assertThat(basedir).as("Maven surefire/failsafe should set system property 'basedir'").isNotBlank(); + return Paths.get(basedir).resolve("../packaging/target/arthas-bin").normalize(); + } + + private static int findFreePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + socket.setReuseAddress(true); + return socket.getLocalPort(); + } + } + + private static void waitForPortOpen(String host, int port, Duration timeout) throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (System.nanoTime() < deadline) { + try (Socket socket = new Socket()) { + socket.connect(new InetSocketAddress(host, port), 500); + return; + } catch (IOException ignored) { + Thread.sleep(200); + } + } + throw new IllegalStateException("等待端口监听超时: " + host + ":" + port); + } + + private static boolean isWindows() { + String os = System.getProperty("os.name"); + return os != null && os.toLowerCase(Locale.ROOT).contains("win"); + } + + private static void deleteDirectoryQuietly(Path dir) { + try { + if (!Files.exists(dir)) { + return; + } + Files.walk(dir) + .sorted(Comparator.reverseOrder()) + .map(Path::toFile) + .forEach(File::delete); + } catch (Exception ignored) { + } + } + + private static final class ProcessPid { + private ProcessPid() { + } + + static long pidOf(Process process) { + // Java 9+ + try { + return (Long) Process.class.getMethod("pid").invoke(process); + } catch (Exception ignored) { + // Java 8 fallback + } + try { + java.lang.reflect.Field pidField = process.getClass().getDeclaredField("pid"); + pidField.setAccessible(true); + Object value = pidField.get(process); + if (value instanceof Number) { + return ((Number) value).longValue(); + } + throw new IllegalStateException("Unsupported pid field type: " + value); + } catch (Exception e) { + throw new IllegalStateException("无法获取目标 JVM pid", e); + } + } + } + + private static final class StreamableMcpHttpClient { + private final String baseUrl; + private final String mcpEndpoint; + + StreamableMcpHttpClient(String host, int port, String mcpEndpoint) { + this.baseUrl = "http://" + host + ":" + port; + this.mcpEndpoint = mcpEndpoint; + } + + String initialize() throws Exception { + McpSchema.InitializeRequest init = new McpSchema.InitializeRequest( + McpSchema.LATEST_PROTOCOL_VERSION, + new McpSchema.ClientCapabilities(null, null, null, null), + new McpSchema.Implementation("arthas-mcp-it", "1.0.0") + ); + McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest( + McpSchema.JSONRPC_VERSION, + McpSchema.METHOD_INITIALIZE, + 1, + init + ); + + HttpURLConnection conn = openPostConnection(null); + writeJson(conn, request); + + int code = conn.getResponseCode(); + String body = readBody(conn); + if (code != 200) { + throw new IllegalStateException("initialize 失败: http=" + code + ", body=" + body); + } + + String sessionId = conn.getHeaderField(HttpHeaders.MCP_SESSION_ID); + if (sessionId == null || sessionId.trim().isEmpty()) { + throw new IllegalStateException("initialize 未返回 mcp-session-id header, body=" + body); + } + + McpSchema.JSONRPCMessage msg = McpSchema.deserializeJsonRpcMessage(OBJECT_MAPPER, body); + if (!(msg instanceof McpSchema.JSONRPCResponse)) { + throw new IllegalStateException("initialize 响应不是 JSONRPCResponse: " + body); + } + McpSchema.JSONRPCResponse resp = (McpSchema.JSONRPCResponse) msg; + if (resp.getError() != null) { + throw new IllegalStateException("initialize 返回 error: " + OBJECT_MAPPER.writeValueAsString(resp.getError())); + } + return sessionId; + } + + void sendInitializedNotification(String sessionId) throws Exception { + McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification( + McpSchema.JSONRPC_VERSION, + McpSchema.METHOD_NOTIFICATION_INITIALIZED, + Collections.emptyMap() + ); + HttpURLConnection conn = openPostConnection(sessionId); + writeJson(conn, notification); + int code = conn.getResponseCode(); + if (code != 202 && code != 200) { + throw new IllegalStateException("notifications/initialized 失败: http=" + code + ", body=" + readBody(conn)); + } + } + + McpSchema.ListToolsResult listTools(String sessionId) throws Exception { + McpSchema.PaginatedRequest params = new McpSchema.PaginatedRequest(null); + McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest( + McpSchema.JSONRPC_VERSION, + McpSchema.METHOD_TOOLS_LIST, + 2, + params + ); + McpSchema.JSONRPCResponse response = postRequestExpectSseResponse(sessionId, request, Duration.ofSeconds(30)); + if (response.getError() != null) { + throw new IllegalStateException("tools/list 返回 error: " + OBJECT_MAPPER.writeValueAsString(response.getError())); + } + return OBJECT_MAPPER.convertValue(response.getResult(), McpSchema.ListToolsResult.class); + } + + McpSchema.CallToolResult callTool(String sessionId, String toolName, Map arguments) throws Exception { + McpSchema.CallToolRequest params = new McpSchema.CallToolRequest(toolName, arguments, null); + McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest( + McpSchema.JSONRPC_VERSION, + McpSchema.METHOD_TOOLS_CALL, + 3, + params + ); + McpSchema.JSONRPCResponse response = postRequestExpectSseResponse(sessionId, request, Duration.ofSeconds(60)); + if (response.getError() != null) { + throw new IllegalStateException("tools/call 返回 error: " + OBJECT_MAPPER.writeValueAsString(response.getError())); + } + return OBJECT_MAPPER.convertValue(response.getResult(), McpSchema.CallToolResult.class); + } + + private HttpURLConnection openPostConnection(String sessionId) throws IOException { + URL url = new URL(baseUrl + mcpEndpoint); + HttpURLConnection conn = (HttpURLConnection) url.openConnection(); + conn.setRequestMethod("POST"); + conn.setConnectTimeout(5_000); + conn.setReadTimeout(30_000); + conn.setDoOutput(true); + conn.setRequestProperty("Content-Type", "application/json"); + conn.setRequestProperty("Accept", "text/event-stream, application/json"); + if (sessionId != null) { + conn.setRequestProperty(HttpHeaders.MCP_SESSION_ID, sessionId); + } + return conn; + } + + private static void writeJson(HttpURLConnection conn, Object body) throws IOException { + byte[] bytes = OBJECT_MAPPER.writeValueAsBytes(body); + conn.setFixedLengthStreamingMode(bytes.length); + try (OutputStream os = conn.getOutputStream()) { + os.write(bytes); + } + } + + private static String readBody(HttpURLConnection conn) throws IOException { + InputStream is = null; + try { + is = conn.getInputStream(); + } catch (IOException e) { + is = conn.getErrorStream(); + } + if (is == null) { + return ""; + } + try (BufferedReader reader = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) { + StringBuilder sb = new StringBuilder(); + String line; + while ((line = reader.readLine()) != null) { + sb.append(line); + } + return sb.toString(); + } + } + + private McpSchema.JSONRPCResponse postRequestExpectSseResponse(String sessionId, McpSchema.JSONRPCRequest request, Duration timeout) + throws Exception { + HttpURLConnection conn = openPostConnection(sessionId); + conn.setReadTimeout((int) timeout.toMillis()); + writeJson(conn, request); + + int code = conn.getResponseCode(); + if (code != 200) { + throw new IllegalStateException(request.getMethod() + " 失败: http=" + code + ", body=" + readBody(conn)); + } + + try (InputStream is = conn.getInputStream()) { + return readJsonRpcResponseFromSse(is, request.getId()); + } + } + + private static McpSchema.JSONRPCResponse readJsonRpcResponseFromSse(InputStream inputStream, Object expectedId) throws Exception { + BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)); + String line; + String data = null; + while ((line = reader.readLine()) != null) { + if (line.startsWith("data:")) { + data = line.substring("data:".length()).trim(); + continue; + } + if (line.isEmpty() && data != null) { + McpSchema.JSONRPCMessage msg = McpSchema.deserializeJsonRpcMessage(OBJECT_MAPPER, data); + data = null; + if (msg instanceof McpSchema.JSONRPCResponse) { + McpSchema.JSONRPCResponse resp = (McpSchema.JSONRPCResponse) msg; + if (Objects.equals(resp.getId(), expectedId)) { + return resp; + } + } + } + } + throw new IllegalStateException("未从 SSE 流中读取到期望的 JSONRPCResponse, id=" + expectedId); + } + } +} diff --git a/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/TargetJvmApp.java b/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/TargetJvmApp.java new file mode 100644 index 000000000..ce99c7aa8 --- /dev/null +++ b/arthas-mcp-integration-test/src/test/java/com/taobao/arthas/mcp/it/TargetJvmApp.java @@ -0,0 +1,25 @@ +package com.taobao.arthas.mcp.it; + +public class TargetJvmApp { + + private static final TargetJvmApp INSTANCE = new TargetJvmApp(); + + /** + * 持续调用的方法，用于触发 watch/trace/monitor/stack/tt 等需要方法执行事件的工具。 + */ + public int hotMethod(int value) { + return compute(value) + 1; + } + + private int compute(int value) { + return value * 2; + } + + public static void main(String[] args) throws Exception { + System.out.println("TargetJvmApp started."); + while (true) { + INSTANCE.hotMethod((int) (System.nanoTime() & 0xFF)); + Thread.sleep(50); + } + } +} diff --git a/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/AbstractArthasTool.java b/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/AbstractArthasTool.java index 2817ee409..09ea3cbce 100644 --- a/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/AbstractArthasTool.java +++ b/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/AbstractArthasTool.java @@ -22,6 +22,9 @@ public abstract class AbstractArthasTool { protected final Logger logger = LoggerFactory.getLogger(this.getClass()); public static final int DEFAULT_TIMEOUT_SECONDS = (int) (StreamableToolUtils.DEFAULT_TIMEOUT_MS / 1000); + + private static final long DEFAULT_ASYNC_START_RETRY_INTERVAL_MS = 100L; + private static final long DEFAULT_ASYNC_START_MAX_WAIT_MS = 3000L; /** * 工具执行上下文，包含所有必要的上下文信息 @@ -112,8 +115,9 @@ public abstract class AbstractArthasTool { Integer expectedResultCount, Integer pollIntervalMs, Integer timeoutMs, String successMessage) { + ToolExecutionContext execContext = null; try { - ToolExecutionContext execContext = new ToolExecutionContext(toolContext, true); + execContext = new ToolExecutionContext(toolContext, true); logger.info("Starting streamable execution: {}", commandStr); @@ -122,7 +126,11 @@ public abstract class AbstractArthasTool { execContext.getCommandContext().setSessionUserId(execContext.getUserId()); } - Map asyncResult = execContext.getCommandContext().executeAsync(commandStr); + Map asyncResult = executeAsyncWithRetry(execContext, commandStr, timeoutMs); + if (!isAsyncExecutionStarted(asyncResult)) { + String errorMessage = asyncResult != null ? String.valueOf(asyncResult.get("error")) : "unknown error"; + return JsonParser.toJson(createErrorResponse("Failed to start command: " + errorMessage)); + } logger.debug("Async execution started: {}", asyncResult); Map results = executeAndCollectResults( @@ -154,9 +162,76 @@ public abstract class AbstractArthasTool { } catch (Exception e) { logger.error("Error executing streamable command: {}", commandStr, e); return JsonParser.toJson(createErrorResponse("Error executing command: " + e.getMessage())); + } finally { + if (execContext != null) { + try { + // 确保前台任务被及时释放，避免占用 session 影响后续 streamable 工具执行 + execContext.getCommandContext().interruptJob(); + } catch (Exception ignored) { + } + } } } + private static boolean isAsyncExecutionStarted(Map asyncResult) { + if (asyncResult == null) { + return false; + } + Object success = asyncResult.get("success"); + return Boolean.TRUE.equals(success); + } + + private static boolean isRetryableAsyncStartError(Map asyncResult) { + if (asyncResult == null) { + return false; + } + Object success = asyncResult.get("success"); + if (Boolean.TRUE.equals(success)) { + return false; + } + Object error = asyncResult.get("error"); + if (error == null) { + return false; + } + String message = String.valueOf(error); + return message.contains("Another job is running") || message.contains("Another command is executing"); + } + + private static Map executeAsyncWithRetry(ToolExecutionContext execContext, String commandStr, Integer timeoutMs) { + long maxWaitMs = DEFAULT_ASYNC_START_MAX_WAIT_MS; + if (timeoutMs != null && timeoutMs > 0) { + maxWaitMs = Math.min(maxWaitMs, timeoutMs); + } + + long deadline = System.currentTimeMillis() + maxWaitMs; + Map asyncResult = null; + + while (System.currentTimeMillis() < deadline) { + asyncResult = execContext.getCommandContext().executeAsync(commandStr); + if (isAsyncExecutionStarted(asyncResult)) { + return asyncResult; + } + + if (isRetryableAsyncStartError(asyncResult)) { + try { + execContext.getCommandContext().interruptJob(); + } catch (Exception ignored) { + } + try { + Thread.sleep(DEFAULT_ASYNC_START_RETRY_INTERVAL_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return asyncResult; + } + continue; + } + + return asyncResult; + } + + return asyncResult; + } + protected StringBuilder buildCommand(String baseCommand) { return new StringBuilder(baseCommand); } diff --git a/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/StreamableToolUtils.java b/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/StreamableToolUtils.java index 14b05d4bc..b14cff52b 100644 --- a/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/StreamableToolUtils.java +++ b/core/src/main/java/com/taobao/arthas/core/mcp/tool/function/StreamableToolUtils.java @@ -108,9 +108,13 @@ public final class StreamableToolUtils { // 判断是否应该结束 // 如果是TERMINATED状态，或者命令已完成且允许输入次数大于等于2，或者实际结果数量达到预期结果数量 - if ("TERMINATED".equals(jobStatus) - || (commandCompleted && allowInputCount >= MIN_ALLOW_INPUT_COUNT_TO_COMPLETE) - || (expectedResultCount != null && totalResultCount >= expectedResultCount)) { + boolean hasExpectedResultCount = (expectedResultCount != null); + boolean reachedExpectedResultCount = hasExpectedResultCount && totalResultCount >= expectedResultCount; + boolean allowInputCompletion = !hasExpectedResultCount + && commandCompleted + && allowInputCount >= MIN_ALLOW_INPUT_COUNT_TO_COMPLETE; + + if ("TERMINATED".equals(jobStatus) || allowInputCompletion || reachedExpectedResultCount) { logger.info("Command completed. Total results collected: {}, Expected: {}", totalResultCount, expectedResultCount); break; } diff --git a/pom.xml b/pom.xml index 268a382be..6d4d8c991 100644 --- a/pom.xml +++ b/pom.xml @@ -212,6 +212,12 @@ zt-zip 1.16 + + + io.modelcontextprotocol.sdk + mcp + 0.17.0 + @@ -236,6 +242,16 @@ + + jdk17 + + [17,) + + + arthas-mcp-integration-test + + + full