feat: Add MCP (Model Context Protocol) base protocol implementation with Netty (#3039)
JavaCI / ubuntu_build (11) (push) Has been cancelled Details
JavaCI / ubuntu_build (17) (push) Has been cancelled Details
JavaCI / ubuntu_build (19) (push) Has been cancelled Details
JavaCI / ubuntu_build (8) (push) Has been cancelled Details
JavaCI / windows_build (11) (push) Has been cancelled Details
JavaCI / windows_build (8) (push) Has been cancelled Details
JavaCI / macos_build (11, macos-14) (push) Has been cancelled Details
JavaCI / macos_build (11, macos-latest) (push) Has been cancelled Details
JavaCI / macos_build (8, macos-14) (push) Has been cancelled Details
JavaCI / macos_build (8, macos-latest) (push) Has been cancelled Details

This commit is contained in:
Zhang Yu 2025-07-30 15:50:51 +08:00 committed by GitHub
parent 0aeffc9bda
commit 1522341a5c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
60 changed files with 7149 additions and 4 deletions

View File

@ -263,6 +263,12 @@
<scope>provided</scope> <scope>provided</scope>
<optional>true</optional> <optional>true</optional>
</dependency> </dependency>
<dependency>
<groupId>com.taobao.arthas</groupId>
<artifactId>arthas-mcp-server</artifactId>
<version>${project.version}</version>
</dependency>
</dependencies> </dependencies>
</project> </project>

View File

@ -0,0 +1,319 @@
package com.taobao.arthas.core.command;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONWriter;
import com.taobao.arthas.core.command.model.ResultModel;
import com.taobao.arthas.core.distribution.impl.PackingResultDistributorImpl;
import com.taobao.arthas.core.shell.cli.CliToken;
import com.taobao.arthas.core.shell.cli.CliTokens;
import com.taobao.arthas.core.shell.cli.Completion;
import com.taobao.arthas.core.shell.handlers.Handler;
import com.taobao.arthas.core.shell.session.Session;
import com.taobao.arthas.core.shell.session.SessionManager;
import com.taobao.arthas.core.shell.system.Job;
import com.taobao.arthas.core.shell.system.JobController;
import com.taobao.arthas.core.shell.system.JobListener;
import com.taobao.arthas.core.shell.system.impl.InternalCommandManager;
import com.taobao.arthas.core.shell.term.SignalHandler;
import com.taobao.arthas.core.shell.term.Term;
import com.taobao.arthas.mcp.server.CommandExecutor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
/**
* 命令执行器用于执行Arthas命令
*/
public class CommandExecutorImpl implements CommandExecutor {
private static final Logger logger = LoggerFactory.getLogger(CommandExecutorImpl.class);
private final SessionManager sessionManager;
private final JobController jobController;
private final InternalCommandManager commandManager;
public CommandExecutorImpl(SessionManager sessionManager) {
this.sessionManager = sessionManager;
this.commandManager = sessionManager.getCommandManager();
this.jobController = sessionManager.getJobController();
}
/**
* 同步执行命令
* @param commandLine 命令行
* @param timeout 超时时间(毫秒)
* @return 命令执行结果
*/
@Override
public Map<String, Object> execute(String commandLine, long timeout) {
if (commandLine == null || commandLine.trim().isEmpty()) {
logger.error("Command line is null or empty");
return createErrorResult(commandLine, "Command line is null or empty");
}
List<CliToken> tokens = CliTokens.tokenize(commandLine);
if (tokens.isEmpty()) {
logger.error("No command found in command line: {}", commandLine);
return createErrorResult(commandLine, "No command found in command line");
}
Map<String, Object> result = new TreeMap<>();
Session session = null;
PackingResultDistributorImpl resultDistributor = null;
Job job = null;
try {
session = sessionManager.createSession();
if (session == null) {
logger.error("Failed to create session for command: {}", commandLine);
return createErrorResult(commandLine, "Failed to create session");
}
resultDistributor = new PackingResultDistributorImpl(session);
InMemoryTerm term = new InMemoryTerm();
term.setSession(session);
job = jobController.createJob(commandManager, tokens, session,
new JobHandle(), term, resultDistributor);
if (job == null) {
logger.error("Failed to create job for command: {}", commandLine);
return createErrorResult(commandLine, "Failed to create job");
}
job.run();
boolean finished = waitForJob(job, (int) timeout);
if (!finished) {
logger.warn("Command timeout after {} ms: {}", timeout, commandLine);
job.interrupt();
return createTimeoutResult(commandLine, timeout);
}
result.put("command", commandLine);
result.put("success", true);
result.put("executionTime", System.currentTimeMillis());
List<ResultModel> results = resultDistributor.getResults();
if (results != null && !results.isEmpty()) {
result.put("results", results);
result.put("resultCount", results.size());
} else {
result.put("results", results);
result.put("resultCount", 0);
}
String termOutput = term.getOutput();
if (termOutput != null && !termOutput.trim().isEmpty()) {
result.put("terminalOutput", termOutput);
}
logger.info("Command executed successfully: {}", commandLine);
return result;
} catch (Exception e) {
logger.error("Error executing command: {}", commandLine, e);
return createErrorResult(commandLine, "Error executing command: " + e.getMessage());
} finally {
if (session != null) {
try {
sessionManager.removeSession(session.getSessionId());
} catch (Exception e) {
logger.warn("Error removing session", e);
}
}
}
}
private boolean waitForJob(Job job, int timeout) {
long startTime = System.currentTimeMillis();
while (true) {
switch (job.status()) {
case STOPPED:
case TERMINATED:
return true;
}
if (System.currentTimeMillis() - startTime > timeout) {
return false;
}
try {
Thread.sleep(100);
} catch (InterruptedException e) {
}
}
}
private Map<String, Object> createErrorResult(String commandLine, String errorMessage) {
Map<String, Object> result = new TreeMap<>();
result.put("command", commandLine);
result.put("success", false);
result.put("error", errorMessage);
result.put("executionTime", System.currentTimeMillis());
return result;
}
/**
* 创建超时结果
*/
private Map<String, Object> createTimeoutResult(String commandLine, long timeout) {
Map<String, Object> result = new TreeMap<>();
result.put("command", commandLine);
result.put("success", false);
result.put("error", "Command timeout after " + timeout + " ms");
result.put("timeout", true);
result.put("executionTime", System.currentTimeMillis());
return result;
}
private static class JobHandle implements JobListener {
private final CountDownLatch latch = new CountDownLatch(1);
@Override
public void onForeground(Job job) {
}
@Override
public void onBackground(Job job) {
}
@Override
public void onTerminated(Job job) {
latch.countDown();
}
@Override
public void onSuspend(Job job) {
}
public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
return latch.await(timeout, unit);
}
}
public static class InMemoryTerm implements Term {
private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
private Session session;
private volatile boolean closed = false;
@Override
public Term setSession(Session session) {
this.session = session;
return this;
}
public Session getSession() {
return session;
}
@Override
public void readline(String prompt, Handler<String> lineHandler) {
}
@Override
public void readline(String prompt, Handler<String> lineHandler, Handler<Completion> completionHandler) {
}
@Override
public Term closeHandler(Handler<Void> handler) {
return this;
}
@Override
public long lastAccessedTime() {
return System.currentTimeMillis();
}
@Override
public String type() {
return "inmemory";
}
@Override
public int width() {
return 120;
}
@Override
public int height() {
return 40;
}
@Override
public Term resizehandler(Handler<Void> handler) {
return this;
}
@Override
public Term stdinHandler(Handler<String> handler) {
return this;
}
@Override
public Term stdoutHandler(io.termd.core.function.Function<String, String> handler) {
return this;
}
@Override
public synchronized Term write(String data) {
if (closed) {
return this;
}
try {
if (data != null) {
outputStream.write(data.getBytes(StandardCharsets.UTF_8));
}
} catch (Exception e) {
logger.error("Error writing to terminal", e);
}
return this;
}
@Override
public Term interruptHandler(SignalHandler handler) {
return this;
}
@Override
public Term suspendHandler(SignalHandler handler) {
return this;
}
@Override
public synchronized void close() {
closed = true;
try {
outputStream.close();
} catch (Exception e) {
logger.error("Error closing output stream", e);
}
}
@Override
public Term echo(String text) {
return write(text);
}
public synchronized String getOutput() {
try {
return outputStream.toString(StandardCharsets.UTF_8.name());
} catch (UnsupportedEncodingException e) {
logger.error("Error getting output", e);
return "";
}
}
public synchronized void clearOutput() {
outputStream.reset();
}
public boolean isClosed() {
return closed;
}
}
}

View File

@ -43,6 +43,7 @@ import com.taobao.arthas.common.SocketUtils;
import com.taobao.arthas.core.advisor.Enhancer; import com.taobao.arthas.core.advisor.Enhancer;
import com.taobao.arthas.core.advisor.TransformerManager; import com.taobao.arthas.core.advisor.TransformerManager;
import com.taobao.arthas.core.command.BuiltinCommandPack; import com.taobao.arthas.core.command.BuiltinCommandPack;
import com.taobao.arthas.core.command.CommandExecutorImpl;
import com.taobao.arthas.core.command.view.ResultViewResolver; import com.taobao.arthas.core.command.view.ResultViewResolver;
import com.taobao.arthas.core.config.BinderUtils; import com.taobao.arthas.core.config.BinderUtils;
import com.taobao.arthas.core.config.Configure; import com.taobao.arthas.core.config.Configure;
@ -77,6 +78,10 @@ import com.taobao.arthas.core.util.UserStatUtil;
import com.taobao.arthas.core.util.affect.EnhancerAffect; import com.taobao.arthas.core.util.affect.EnhancerAffect;
import com.taobao.arthas.core.util.matcher.WildcardMatcher; import com.taobao.arthas.core.util.matcher.WildcardMatcher;
import com.taobao.arthas.mcp.server.ArthasMcpBootstrap;
import com.taobao.arthas.mcp.server.CommandExecutor;
import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerTransportProvider;
import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelFuture;
import io.netty.channel.nio.NioEventLoopGroup; import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.util.concurrent.DefaultThreadFactory; import io.netty.util.concurrent.DefaultThreadFactory;
@ -125,6 +130,8 @@ public class ArthasBootstrap {
private HttpApiHandler httpApiHandler; private HttpApiHandler httpApiHandler;
private McpRequestHandler mcpRequestHandler;
private HttpSessionManager httpSessionManager; private HttpSessionManager httpSessionManager;
private SecurityAuthenticator securityAuthenticator; private SecurityAuthenticator securityAuthenticator;
@ -462,6 +469,11 @@ public class ArthasBootstrap {
//http api handler //http api handler
httpApiHandler = new HttpApiHandler(historyManager, sessionManager); httpApiHandler = new HttpApiHandler(historyManager, sessionManager);
// Mcp Server
CommandExecutor commandExecutor = new CommandExecutorImpl(sessionManager);
ArthasMcpBootstrap arthasMcpBootstrap = new ArthasMcpBootstrap(commandExecutor);
this.mcpRequestHandler = arthasMcpBootstrap.start().getMcpRequestHandler();
logger().info("as-server listening on network={};telnet={};http={};timeout={};", configure.getIp(), logger().info("as-server listening on network={};telnet={};http={};timeout={};", configure.getIp(),
configure.getTelnetPort(), configure.getHttpPort(), options.getConnectionTimeout()); configure.getTelnetPort(), configure.getHttpPort(), options.getConnectionTimeout());
@ -671,6 +683,10 @@ public class ArthasBootstrap {
return httpApiHandler; return httpApiHandler;
} }
public McpRequestHandler getMcpRequestHandler() {
return mcpRequestHandler;
}
public File getOutputPath() { public File getOutputPath() {
return outputPath; return outputPath;
} }

View File

@ -5,6 +5,7 @@ import com.alibaba.arthas.deps.org.slf4j.LoggerFactory;
import com.taobao.arthas.common.IOUtils; import com.taobao.arthas.common.IOUtils;
import com.taobao.arthas.core.server.ArthasBootstrap; import com.taobao.arthas.core.server.ArthasBootstrap;
import com.taobao.arthas.core.shell.term.impl.http.api.HttpApiHandler; import com.taobao.arthas.core.shell.term.impl.http.api.HttpApiHandler;
import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler;
import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener; import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelHandlerContext;
@ -37,6 +38,7 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
private HttpApiHandler httpApiHandler; private HttpApiHandler httpApiHandler;
private McpRequestHandler mcpRequestHandler;
public HttpRequestHandler(String wsUri) { public HttpRequestHandler(String wsUri) {
this(wsUri, ArthasBootstrap.getInstance().getOutputPath()); this(wsUri, ArthasBootstrap.getInstance().getOutputPath());
@ -47,6 +49,7 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
this.dir = dir; this.dir = dir;
dir.mkdirs(); dir.mkdirs();
this.httpApiHandler = ArthasBootstrap.getInstance().getHttpApiHandler(); this.httpApiHandler = ArthasBootstrap.getInstance().getHttpApiHandler();
this.mcpRequestHandler = ArthasBootstrap.getInstance().getMcpRequestHandler();
} }
@Override @Override
@ -65,12 +68,22 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
} }
boolean isFileResponseFinished = false; boolean isFileResponseFinished = false;
boolean isMcpHandled = false;
try { try {
//handle http restful api //handle http restful api
if ("/api".equals(path)) { if ("/api".equals(path)) {
response = httpApiHandler.handle(ctx, request); response = httpApiHandler.handle(ctx, request);
} }
//handle mcp request
String messageEndpoint = mcpRequestHandler.getMessageEndpoint();
String sseEndpoint = mcpRequestHandler.getSseEndpoint();
if (sseEndpoint.equals(path) || messageEndpoint.equals(path)) {
mcpRequestHandler.handle(ctx, request);
isMcpHandled = true;
return;
}
//handle webui requests //handle webui requests
if (path.equals("/ui")) { if (path.equals("/ui")) {
response = createRedirectResponse(request, "/ui/"); response = createRedirectResponse(request, "/ui/");
@ -101,7 +114,7 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
if (response == null) { if (response == null) {
response = createResponse(request, HttpResponseStatus.INTERNAL_SERVER_ERROR, "Server error"); response = createResponse(request, HttpResponseStatus.INTERNAL_SERVER_ERROR, "Server error");
} }
if (!isFileResponseFinished) { if (!isFileResponseFinished && !isMcpHandled) {
ChannelFuture future = writeResponse(ctx, response); ChannelFuture future = writeResponse(ctx, response);
future.addListener(ChannelFutureListener.CLOSE); future.addListener(ChannelFutureListener.CLOSE);
} }
@ -112,10 +125,10 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
private ChannelFuture writeResponse(ChannelHandlerContext ctx, HttpResponse response) { private ChannelFuture writeResponse(ChannelHandlerContext ctx, HttpResponse response) {
// try to add content-length header for DefaultFullHttpResponse // try to add content-length header for DefaultFullHttpResponse
if (!HttpUtil.isTransferEncodingChunked(response) if (!HttpUtil.isTransferEncodingChunked(response)
&& response instanceof DefaultFullHttpResponse) { && response instanceof DefaultFullHttpResponse) {
response.headers().set(HttpHeaderNames.CONNECTION, HttpHeaderValues.CLOSE); response.headers().set(HttpHeaderNames.CONNECTION, HttpHeaderValues.CLOSE);
response.headers().set(HttpHeaderNames.CONTENT_LENGTH, response.headers().set(HttpHeaderNames.CONTENT_LENGTH,
((DefaultFullHttpResponse) response).content().readableBytes()); ((DefaultFullHttpResponse) response).content().readableBytes());
return ctx.writeAndFlush(response); return ctx.writeAndFlush(response);
} }
@ -131,7 +144,7 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
URL res = HttpTtyConnection.class.getResource("/com/taobao/arthas/core/http" + path); URL res = HttpTtyConnection.class.getResource("/com/taobao/arthas/core/http" + path);
if (res != null) { if (res != null) {
fullResp = new DefaultFullHttpResponse(request.protocolVersion(), fullResp = new DefaultFullHttpResponse(request.protocolVersion(),
HttpResponseStatus.OK); HttpResponseStatus.OK);
in = res.openStream(); in = res.openStream();
byte[] tmp = new byte[256]; byte[] tmp = new byte[256];
for (int l = 0; l != -1; l = in.read(tmp)) { for (int l = 0; l != -1; l = in.read(tmp)) {

View File

@ -0,0 +1,22 @@
# arthas-mcp-server
## 项目简介
`arthas-mcp-server` 是 [Arthas](https://github.com/alibaba/arthas) 的实验模块,实现了基于 MCPModel Context Protocol协议的服务端。该模块通过 HTTP/Netty 提供统一的 JSON-RPC 2.0 接口,支持 AI 使用工具调用的方式执行 arthas 的命令。
## 快速开始
正常启动服务之后服务对外暴露8563在 cherry-studio/cline 等 ai 客户端中配置:
在设置中添加 MCP 服务器:
```JSON
{
"mcpServers": {
"arthas-mcp": {
"transport": "sse",
"url": "http://localhost:8563/sse"
}
}
}
```

View File

@ -0,0 +1,93 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<artifactId>arthas-all</artifactId>
<groupId>com.taobao.arthas</groupId>
<version>${revision}</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>arthas-mcp-server</artifactId>
<name>arthas-mcp-server</name>
<url>https://github.com/alibaba/arthas</url>
<properties>
<maven.compiler.source>8</maven.compiler.source>
<maven.compiler.target>8</maven.compiler.target>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
</properties>
<dependencies>
<!-- Netty -->
<dependency>
<groupId>io.netty</groupId>
<artifactId>netty-all</artifactId>
<version>${netty.version}</version>
</dependency>
<!-- Jackson dependencies -->
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-core</artifactId>
<version>2.18.1</version>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-databind</artifactId>
<version>2.18.1</version>
</dependency>
<!-- log dependencies -->
<dependency>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-api</artifactId>
</dependency>
<!-- Test dependencies -->
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.assertj</groupId>
<artifactId>assertj-core</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
<build>
<finalName>arthas-mcp-server</finalName>
<resources>
<resource>
<directory>src/main/resources</directory>
<filtering>false</filtering>
<includes>
<include>**/*</include>
</includes>
</resource>
</resources>
<plugins>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-compiler-plugin</artifactId>
<configuration>
<source>8</source>
<target>8</target>
<encoding>UTF-8</encoding>
<parameters>true</parameters>
</configuration>
</plugin>
</plugins>
</build>
</project>

View File

@ -0,0 +1,49 @@
package com.taobao.arthas.mcp.server;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Arthas MCP Bootstrap class
*
* @author Yeaury
*/
public class ArthasMcpBootstrap {
private static final Logger logger = LoggerFactory.getLogger(ArthasMcpBootstrap.class);
private ArthasMcpServer mcpServer;
private final CommandExecutor commandExecutor;
private static ArthasMcpBootstrap instance;
public ArthasMcpBootstrap(CommandExecutor commandExecutor) {
this.commandExecutor = commandExecutor;
instance = this;
}
public static ArthasMcpBootstrap getInstance() {
return instance;
}
public CommandExecutor getCommandExecutor() {
return commandExecutor;
}
public ArthasMcpServer start() {
try {
// Create and start MCP server
mcpServer = new ArthasMcpServer();
mcpServer.start();
logger.info("Arthas MCP server initialized successfully");
return mcpServer;
} catch (Exception e) {
logger.error("Failed to initialize Arthas MCP server", e);
throw new RuntimeException("Failed to initialize Arthas MCP server", e);
}
}
public void shutdown() {
if (mcpServer != null) {
mcpServer.stop();
}
}
}

View File

@ -0,0 +1,141 @@
package com.taobao.arthas.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.config.McpServerProperties;
import com.taobao.arthas.mcp.server.protocol.server.McpNettyServer;
import com.taobao.arthas.mcp.server.protocol.server.McpServer;
import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler;
import com.taobao.arthas.mcp.server.protocol.server.transport.HttpNettyServerTransportProvider;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.*;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerTransportProvider;
import com.taobao.arthas.mcp.server.tool.DefaultToolCallbackProvider;
import com.taobao.arthas.mcp.server.tool.ToolCallback;
import com.taobao.arthas.mcp.server.tool.ToolCallbackProvider;
import com.taobao.arthas.mcp.server.tool.util.McpToolUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
/**
* Arthas MCP Server
* Used to expose HTTP service after Arthas startup
*/
public class ArthasMcpServer {
private static final Logger logger = LoggerFactory.getLogger(ArthasMcpServer.class);
private McpNettyServer server;
private final int port;
private final String bindAddress;
private McpRequestHandler mcpRequestHandler;
public ArthasMcpServer() {
this(8080, "localhost");
}
public ArthasMcpServer(int port, String bindAddress) {
this.port = port;
this.bindAddress = bindAddress;
}
public McpRequestHandler getMcpRequestHandler() {
return mcpRequestHandler;
}
/**
* Start MCP server
*/
public void start() {
try {
// 1. Create server configuration
McpServerProperties properties = new McpServerProperties.Builder()
.name("arthas-mcp-server")
.version("1.0.0")
.bindAddress(bindAddress)
.port(port)
.messageEndpoint("/sse/message")
.sseEndpoint("/sse")
.toolChangeNotification(true)
.resourceChangeNotification(true)
.promptChangeNotification(true)
.objectMapper(new ObjectMapper())
.build();
// 2. Create transport provider
McpServerTransportProvider transportProvider = createHttpTransportProvider(properties);
mcpRequestHandler = transportProvider.getMcpRequestHandler();
// 3. Create server builder
McpServer.NettySpecification serverBuilder = McpServer.netty(transportProvider)
.serverInfo(new Implementation(properties.getName(), properties.getVersion()))
.capabilities(buildServerCapabilities(properties))
.instructions(properties.getInstructions())
.requestTimeout(properties.getRequestTimeout())
.objectMapper(properties.getObjectMapper() != null ? properties.getObjectMapper() : new ObjectMapper());
ToolCallbackProvider toolCallbackProvider = new DefaultToolCallbackProvider();
ToolCallback[] callbacks = toolCallbackProvider.getToolCallbacks();
List<ToolCallback> providerToolCallbacks = Arrays.stream(callbacks)
.filter(Objects::nonNull)
.collect(Collectors.toList());
serverBuilder.tools(
McpToolUtils.toToolSpecifications(providerToolCallbacks, properties));
server = serverBuilder.build();
logger.info("Arthas MCP server started, listening on {}:{}", bindAddress, port);
} catch (Exception e) {
logger.error("Failed to start Arthas MCP server", e);
throw new RuntimeException("Failed to start Arthas MCP server", e);
}
}
/**
* Create HTTP transport provider
*/
private HttpNettyServerTransportProvider createHttpTransportProvider(McpServerProperties properties) {
return HttpNettyServerTransportProvider.builder()
.messageEndpoint(properties.getMessageEndpoint())
.sseEndpoint(properties.getSseEndpoint())
.objectMapper(properties.getObjectMapper() != null ? properties.getObjectMapper() : new ObjectMapper())
.build();
}
/**
* Build server capabilities configuration
*/
private ServerCapabilities buildServerCapabilities(McpServerProperties properties) {
return ServerCapabilities.builder()
.prompts(new ServerCapabilities.PromptCapabilities(properties.isPromptChangeNotification()))
.resources(new ServerCapabilities.ResourceCapabilities(properties.isResourceSubscribe(), properties.isResourceChangeNotification()))
.tools(new ServerCapabilities.ToolCapabilities(properties.isToolChangeNotification()))
.build();
}
/**
* Stop MCP server
*/
public void stop() {
if (server != null) {
try {
server.closeGracefully().get();
logger.info("Arthas MCP server stopped");
} catch (Exception e) {
logger.error("Failed to stop Arthas MCP server", e);
}
}
}
public static void main(String[] args) {
ArthasMcpServer arthasMcpServer = new ArthasMcpServer();
arthasMcpServer.start();
// Keep the server running
try {
Thread.sleep(Long.MAX_VALUE);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
}

View File

@ -0,0 +1,9 @@
package com.taobao.arthas.mcp.server;
import java.util.Map;
public interface CommandExecutor {
Map<String, Object> execute(String commandLine, long timeout);
}

View File

@ -0,0 +1,289 @@
package com.taobao.arthas.mcp.server.protocol.config;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
/**
* MCP Server Configuration Properties
* Used to manage all configuration items for MCP server.
*
* @author Yeaury
*/
public class McpServerProperties {
/**
* Server basic information
*/
private final String name;
private final String version;
private final String instructions;
/**
* Server capability configuration
*/
private final boolean toolChangeNotification;
private final boolean resourceChangeNotification;
private final boolean promptChangeNotification;
private final boolean resourceSubscribe;
/**
* Transport layer configuration
*/
private final String bindAddress;
private final int port;
private final String messageEndpoint;
private final String sseEndpoint;
/**
* Timeout configuration
*/
private final Duration requestTimeout;
private final Duration initializationTimeout;
private final ObjectMapper objectMapper;
/**
* Private constructor, can only be created through Builder
*/
private McpServerProperties(Builder builder) {
this.name = builder.name;
this.version = builder.version;
this.instructions = builder.instructions;
this.toolChangeNotification = builder.toolChangeNotification;
this.resourceChangeNotification = builder.resourceChangeNotification;
this.promptChangeNotification = builder.promptChangeNotification;
this.resourceSubscribe = builder.resourceSubscribe;
this.bindAddress = builder.bindAddress;
this.port = builder.port;
this.messageEndpoint = builder.messageEndpoint;
this.sseEndpoint = builder.sseEndpoint;
this.requestTimeout = builder.requestTimeout;
this.initializationTimeout = builder.initializationTimeout;
this.objectMapper = builder.objectMapper;
}
/**
* Create Builder with default configuration
*/
public static Builder builder() {
return new Builder();
}
/**
* Get server name
* @return Server name
*/
public String getName() {
return name;
}
/**
* Get server version
* @return Server version
*/
public String getVersion() {
return version;
}
/**
* Get server instructions
* @return Server instructions
*/
public String getInstructions() {
return instructions;
}
/**
* Get tool change notification
* @return Tool change notification
*/
public boolean isToolChangeNotification() {
return toolChangeNotification;
}
/**
* Get resource change notification
* @return Resource change notification
*/
public boolean isResourceChangeNotification() {
return resourceChangeNotification;
}
/**
* Get prompt change notification
* @return Prompt change notification
*/
public boolean isPromptChangeNotification() {
return promptChangeNotification;
}
/**
* Get resource subscribe
* @return Resource subscribe
*/
public boolean isResourceSubscribe() {
return resourceSubscribe;
}
/**
* Get bind address
* @return Bind address
*/
public String getBindAddress() {
return bindAddress;
}
/**
* Get server port
* @return Server port
*/
public int getPort() {
return port;
}
/**
* Get message endpoint
* @return Message endpoint
*/
public String getMessageEndpoint() {
return messageEndpoint;
}
/**
* Get SSE endpoint
* @return SSE endpoint
*/
public String getSseEndpoint() {
return sseEndpoint;
}
/**
* Get request timeout
* @return Request timeout
*/
public Duration getRequestTimeout() {
return requestTimeout;
}
/**
* Get initialization timeout
* @return Initialization timeout
*/
public Duration getInitializationTimeout() {
return initializationTimeout;
}
/**
* Get object mapper
* @return Object mapper
*/
public ObjectMapper getObjectMapper() {
return objectMapper;
}
/**
* Builder class for McpServerProperties
*/
public static class Builder {
// Default values
private String name = "mcp-server";
private String version = "1.0.0";
private String instructions;
private boolean toolChangeNotification = true;
private boolean resourceChangeNotification = true;
private boolean promptChangeNotification = true;
private boolean resourceSubscribe = false;
private String bindAddress = "localhost";
private int port = 8080;
private String messageEndpoint = "/mcp/message";
private String sseEndpoint = "/mcp/sse";
private Duration requestTimeout = Duration.ofSeconds(10);
private Duration initializationTimeout = Duration.ofSeconds(30);
private ObjectMapper objectMapper;
public Builder() {
// Private constructor to prevent direct instantiation
}
public Builder name(String name) {
this.name = name;
return this;
}
public Builder version(String version) {
this.version = version;
return this;
}
public Builder instructions(String instructions) {
this.instructions = instructions;
return this;
}
public Builder toolChangeNotification(boolean toolChangeNotification) {
this.toolChangeNotification = toolChangeNotification;
return this;
}
public Builder resourceChangeNotification(boolean resourceChangeNotification) {
this.resourceChangeNotification = resourceChangeNotification;
return this;
}
public Builder promptChangeNotification(boolean promptChangeNotification) {
this.promptChangeNotification = promptChangeNotification;
return this;
}
public Builder resourceSubscribe(boolean resourceSubscribe) {
this.resourceSubscribe = resourceSubscribe;
return this;
}
public Builder bindAddress(String bindAddress) {
this.bindAddress = bindAddress;
return this;
}
public Builder port(int port) {
this.port = port;
return this;
}
public Builder messageEndpoint(String messageEndpoint) {
this.messageEndpoint = messageEndpoint;
return this;
}
public Builder sseEndpoint(String sseEndpoint) {
this.sseEndpoint = sseEndpoint;
return this;
}
public Builder requestTimeout(Duration requestTimeout) {
this.requestTimeout = requestTimeout;
return this;
}
public Builder initializationTimeout(Duration initializationTimeout) {
this.initializationTimeout = initializationTimeout;
return this;
}
public Builder objectMapper(ObjectMapper objectMapper) {
this.objectMapper = objectMapper;
return this;
}
/**
* Build McpServerProperties instance
*/
public McpServerProperties build() {
return new McpServerProperties(this);
}
}
}

View File

@ -0,0 +1,541 @@
package com.taobao.arthas.mcp.server.protocol.server;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.BiFunction;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.netty.channel.Channel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.taobao.arthas.mcp.server.protocol.spec.*;
import com.taobao.arthas.mcp.server.util.Assert;
import com.taobao.arthas.mcp.server.util.Utils;
/**
* A Netty-based MCP server implementation that provides access to tools, resources, and prompts.
*
* @author Yeaury
*/
public class McpNettyServer {
private static final Logger logger = LoggerFactory.getLogger(McpNettyServer.class);
private final McpServerTransportProvider mcpTransportProvider;
private final ObjectMapper objectMapper;
private final McpSchema.ServerCapabilities serverCapabilities;
private final McpSchema.Implementation serverInfo;
private final String instructions;
private final CopyOnWriteArrayList<McpServerFeatures.ToolSpecification> tools = new CopyOnWriteArrayList<>();
private final CopyOnWriteArrayList<McpSchema.ResourceTemplate> resourceTemplates = new CopyOnWriteArrayList<>();
private final ConcurrentHashMap<String, McpServerFeatures.ResourceSpecification> resources = new ConcurrentHashMap<>();
private final ConcurrentHashMap<String, McpServerFeatures.PromptSpecification> prompts = new ConcurrentHashMap<>();
private McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.DEBUG;
private List<String> protocolVersions = Collections.singletonList(McpSchema.LATEST_PROTOCOL_VERSION);
public McpNettyServer(
McpServerTransportProvider mcpTransportProvider,
ObjectMapper objectMapper,
Duration requestTimeout,
McpServerFeatures.McpServerConfig features) {
this.mcpTransportProvider = mcpTransportProvider;
this.objectMapper = objectMapper;
this.serverInfo = features.getServerInfo();
this.serverCapabilities = features.getServerCapabilities();
this.instructions = features.getInstructions();
this.tools.addAll(features.getTools());
this.resources.putAll(features.getResources());
this.resourceTemplates.addAll(features.getResourceTemplates());
this.prompts.putAll(features.getPrompts());
Map<String, McpServerSession.RequestHandler<?>> requestHandlers = new HashMap<>();
// Initialize request handlers for standard MCP methods
// Ping MUST respond with an empty data, but not NULL response.
requestHandlers.put(McpSchema.METHOD_PING,
(exchange, params) -> CompletableFuture.completedFuture(Collections.emptyMap()));
// Add tools API handlers if the tool capability is enabled
if (this.serverCapabilities.getTools() != null) {
requestHandlers.put(McpSchema.METHOD_TOOLS_LIST, toolsListRequestHandler());
requestHandlers.put(McpSchema.METHOD_TOOLS_CALL, toolsCallRequestHandler());
}
// Add resources API handlers if provided
if (this.serverCapabilities.getResources() != null) {
requestHandlers.put(McpSchema.METHOD_RESOURCES_LIST, resourcesListRequestHandler());
requestHandlers.put(McpSchema.METHOD_RESOURCES_READ, resourcesReadRequestHandler());
requestHandlers.put(McpSchema.METHOD_RESOURCES_TEMPLATES_LIST, resourceTemplateListRequestHandler());
}
// Add prompts API handlers if provider exists
if (this.serverCapabilities.getPrompts() != null) {
requestHandlers.put(McpSchema.METHOD_PROMPT_LIST, promptsListRequestHandler());
requestHandlers.put(McpSchema.METHOD_PROMPT_GET, promptsGetRequestHandler());
}
if (this.serverCapabilities.getLogging() != null) {
requestHandlers.put(McpSchema.METHOD_LOGGING_SET_LEVEL, setLoggerRequestHandler());
}
Map<String, McpServerSession.NotificationHandler> notificationHandlers = new HashMap<>();
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED,
(exchange, params) -> CompletableFuture.completedFuture(null));
List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> rootsChangeConsumers = features
.getRootsChangeConsumers();
if (Utils.isEmpty(rootsChangeConsumers)) {
rootsChangeConsumers = Collections.singletonList(
(exchange, roots) -> CompletableFuture.runAsync(() ->
logger.warn("Roots list changed notification, but no consumers provided. Roots list changed: {}", roots))
);
}
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED,
rootsListChangedNotificationHandler(rootsChangeConsumers));
mcpTransportProvider.setSessionFactory(transport -> {
Channel channel = transport.getChannel();
return new McpServerSession(UUID.randomUUID().toString(),
requestTimeout, transport, this::initializeRequestHandler,
() -> CompletableFuture.completedFuture(null), requestHandlers, notificationHandlers,
channel);
});
}
// ---------------------------------------
// Lifecycle Management
// ---------------------------------------
private CompletableFuture<McpSchema.InitializeResult> initializeRequestHandler(
McpSchema.InitializeRequest initializeRequest) {
return CompletableFuture.supplyAsync(() -> {
logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}",
initializeRequest.getProtocolVersion(), initializeRequest.getCapabilities(),
initializeRequest.getClientInfo());
// The server MUST respond with the highest protocol version it supports
// if
// it does not support the requested (e.g. Client) version.
String serverProtocolVersion = protocolVersions.get(protocolVersions.size() - 1);
if (protocolVersions.contains(initializeRequest.getProtocolVersion())) {
serverProtocolVersion = initializeRequest.getProtocolVersion();
}
else {
logger.warn(
"Client requested unsupported protocol version: {}, " + "so the server will suggest {} instead",
initializeRequest.getProtocolVersion(), serverProtocolVersion);
}
return new McpSchema.InitializeResult(serverProtocolVersion, serverCapabilities, serverInfo, instructions);
});
}
public McpSchema.ServerCapabilities getServerCapabilities() {
return this.serverCapabilities;
}
public McpSchema.Implementation getServerInfo() {
return this.serverInfo;
}
public CompletableFuture<Void> closeGracefully() {
return this.mcpTransportProvider.closeGracefully();
}
public void close() {
this.mcpTransportProvider.close();
}
private McpServerSession.NotificationHandler rootsListChangedNotificationHandler(
List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> rootsChangeConsumers) {
return (exchange, params) -> {
CompletableFuture<McpSchema.ListRootsResult> futureRoots = exchange.listRoots();
return futureRoots.thenCompose(listRootsResult -> {
List<McpSchema.Root> roots = listRootsResult.getRoots();
List<CompletableFuture<?>> futures = new ArrayList<>();
for (BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>> consumer : rootsChangeConsumers) {
CompletableFuture<Void> future = consumer.apply(exchange, roots).exceptionally(error -> {
logger.error("Error handling roots list change notification", error);
return null;
});
futures.add(future);
}
return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
});
};
}
// ---------------------------------------
// Tool Management
// ---------------------------------------
public CompletableFuture<Void> addTool(McpServerFeatures.ToolSpecification toolSpecification) {
if (toolSpecification == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Tool specification must not be null"));
return future;
}
if (toolSpecification.getTool() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Tool must not be null"));
return future;
}
if (toolSpecification.getCall() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Tool call handler must not be null"));
return future;
}
if (this.serverCapabilities.getTools() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Server must be configured with tool capabilities"));
return future;
}
return CompletableFuture.supplyAsync(() -> {
// Check for duplicate tool names
if (this.tools.stream().anyMatch(th -> th.getTool().getName().equals(toolSpecification.getTool().getName()))) {
throw new CompletionException(
new McpError("Tool with name '" + toolSpecification.getTool().getName() + "' already exists"));
}
this.tools.add(toolSpecification);
logger.debug("Added tool handler: {}", toolSpecification.getTool().getName());
return null;
}).thenCompose(ignored -> {
if (this.serverCapabilities.getTools().getListChanged()) {
return notifyToolsListChanged();
}
return CompletableFuture.completedFuture(null);
}).exceptionally(ex -> {
Throwable cause = ex instanceof CompletionException ? ex.getCause() : ex;
logger.error("Error while adding tool", cause);
throw new CompletionException(cause);
});
}
public CompletableFuture<Void> removeTool(String toolName) {
if (toolName == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Tool name must not be null"));
return future;
}
if (this.serverCapabilities.getTools() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Server must be configured with tool capabilities"));
return future;
}
return CompletableFuture.supplyAsync(() -> {
boolean removed = this.tools.removeIf(spec -> spec.getTool().getName().equals(toolName));
if (!removed) {
throw new CompletionException(new McpError("Tool with name '" + toolName + "' not found"));
}
logger.debug("Removed tool handler: {}", toolName);
return null;
}).thenCompose(ignored -> {
if (this.serverCapabilities.getTools().getListChanged()) {
return notifyToolsListChanged();
}
return CompletableFuture.completedFuture(null);
}).exceptionally(ex -> {
Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex;
logger.error("Error while removing tool '{}'", toolName, cause);
throw new CompletionException(cause);
});
}
public CompletableFuture<Void> notifyToolsListChanged() {
logger.debug("Notifying clients about tool list changes");
return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED, null);
}
private McpServerSession.RequestHandler<McpSchema.ListToolsResult> toolsListRequestHandler() {
return (exchange, params) -> {
List<McpSchema.Tool> tools = new ArrayList<>();
for (McpServerFeatures.ToolSpecification toolSpec : this.tools) {
tools.add(toolSpec.getTool());
}
return CompletableFuture.completedFuture(new McpSchema.ListToolsResult(tools, null));
};
}
private McpServerSession.RequestHandler<McpSchema.CallToolResult> toolsCallRequestHandler() {
return (exchange, params) -> {
McpSchema.CallToolRequest callToolRequest = objectMapper.convertValue(params,
new TypeReference<McpSchema.CallToolRequest>() {
});
Optional<McpServerFeatures.ToolSpecification> toolSpecification = this.tools.stream()
.filter(tr -> callToolRequest.getName().equals(tr.getTool().getName()))
.findAny();
if (!toolSpecification.isPresent()) {
CompletableFuture<McpSchema.CallToolResult> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("no tool found: " + callToolRequest.getName()));
return future;
}
return toolSpecification.get().getCall().apply(exchange, callToolRequest.getArguments());
};
}
// ---------------------------------------
// Resource Management
// ---------------------------------------
public CompletableFuture<Void> addResource(McpServerFeatures.ResourceSpecification resourceSpecification) {
if (resourceSpecification == null || resourceSpecification.getResource() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Resource must not be null"));
return future;
}
if (this.serverCapabilities.getResources() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Server must be configured with resource capabilities"));
return future;
}
return CompletableFuture.supplyAsync(() -> {
if (this.resources.putIfAbsent(resourceSpecification.getResource().getUri(), resourceSpecification) != null) {
throw new CompletionException(new McpError(
"Resource with URI '" + resourceSpecification.getResource().getUri() + "' already exists"));
}
logger.debug("Added resource handler: {}", resourceSpecification.getResource().getUri());
return null;
}).thenCompose(ignored -> {
if (this.serverCapabilities.getResources().getListChanged()) {
return notifyResourcesListChanged();
}
return CompletableFuture.completedFuture(null);
}).exceptionally(ex -> {
Throwable cause = ex instanceof CompletionException ? ex.getCause() : ex;
logger.error("Error while adding resource '{}'", resourceSpecification.getResource().getUri(), cause);
throw new CompletionException(cause);
});
}
public CompletableFuture<Void> removeResource(String resourceUri) {
if (resourceUri == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Resource URI must not be null"));
return future;
}
if (this.serverCapabilities.getResources() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Server must be configured with resource capabilities"));
return future;
}
return CompletableFuture.supplyAsync(() -> {
McpServerFeatures.ResourceSpecification removed = this.resources.remove(resourceUri);
if (removed == null) {
throw new CompletionException(new McpError("Resource with URI '" + resourceUri + "' not found"));
}
logger.debug("Removed resource handler: {}", resourceUri);
return null;
}).thenCompose(ignored -> {
if (this.serverCapabilities.getResources().getListChanged()) {
return notifyResourcesListChanged();
}
return CompletableFuture.completedFuture(null);
}).exceptionally(ex -> {
Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex;
logger.error("Error while removing resource '{}'", resourceUri, cause);
throw new CompletionException(cause);
});
}
public CompletableFuture<Void> notifyResourcesListChanged() {
return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null);
}
private McpServerSession.RequestHandler<McpSchema.ListResourcesResult> resourcesListRequestHandler() {
return (exchange, params) -> {
List<McpSchema.Resource> resourceList = new ArrayList<>();
for (McpServerFeatures.ResourceSpecification spec : this.resources.values()) {
resourceList.add(spec.getResource());
}
return CompletableFuture.completedFuture(new McpSchema.ListResourcesResult(resourceList, null));
};
}
private McpServerSession.RequestHandler<McpSchema.ListResourceTemplatesResult> resourceTemplateListRequestHandler() {
return (exchange, params) -> CompletableFuture
.completedFuture(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
}
private McpServerSession.RequestHandler<McpSchema.ReadResourceResult> resourcesReadRequestHandler() {
return (exchange, params) -> {
McpSchema.ReadResourceRequest resourceRequest = objectMapper.convertValue(params,
new TypeReference<McpSchema.ReadResourceRequest>() {
});
String resourceUri = resourceRequest.getUri();
McpServerFeatures.ResourceSpecification specification = this.resources.get(resourceUri);
if (specification != null) {
return specification.getReadHandler().apply(exchange, resourceRequest);
}
CompletableFuture<McpSchema.ReadResourceResult> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Resource not found: " + resourceUri));
return future;
};
}
// ---------------------------------------
// Prompt Management
// ---------------------------------------
public CompletableFuture<Void> addPrompt(McpServerFeatures.PromptSpecification promptSpecification) {
if (promptSpecification == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Prompt specification must not be null"));
return future;
}
if (this.serverCapabilities.getPrompts() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Server must be configured with prompt capabilities"));
return future;
}
return CompletableFuture.supplyAsync(() -> {
McpServerFeatures.PromptSpecification existing = this.prompts
.putIfAbsent(promptSpecification.getPrompt().getName(), promptSpecification);
if (existing != null) {
throw new CompletionException(
new McpError("Prompt with name '" + promptSpecification.getPrompt().getName() + "' already exists"));
}
logger.debug("Added prompt handler: {}", promptSpecification.getPrompt().getName());
return null;
}).thenCompose(ignored -> {
if (this.serverCapabilities.getPrompts().getListChanged()) {
return notifyPromptsListChanged();
}
return CompletableFuture.completedFuture(null);
}).exceptionally(ex -> {
Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex;
logger.error("Error while adding prompt '{}'", promptSpecification.getPrompt().getName(), cause);
throw new CompletionException(cause);
});
}
public CompletableFuture<Void> removePrompt(String promptName) {
if (promptName == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Prompt name must not be null"));
return future;
}
if (this.serverCapabilities.getPrompts() == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Server must be configured with prompt capabilities"));
return future;
}
return CompletableFuture.supplyAsync(() -> {
McpServerFeatures.PromptSpecification removed = this.prompts.remove(promptName);
if (removed == null) {
throw new CompletionException(new McpError("Prompt with name '" + promptName + "' not found"));
}
logger.debug("Removed prompt handler: {}", promptName);
return null;
}).thenCompose(ignored -> {
if (this.serverCapabilities.getPrompts().getListChanged()) {
return notifyPromptsListChanged();
}
return CompletableFuture.completedFuture(null);
}).exceptionally(ex -> {
Throwable cause = ex instanceof CompletionException ? ex.getCause() : ex;
logger.error("Error while removing prompt '{}'", promptName, cause);
throw new CompletionException(cause);
});
}
public CompletableFuture<Void> notifyPromptsListChanged() {
return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_PROMPTS_LIST_CHANGED, null);
}
private McpServerSession.RequestHandler<McpSchema.ListPromptsResult> promptsListRequestHandler() {
return (exchange, params) -> {
List<McpSchema.Prompt> promptList = new ArrayList<>();
for (McpServerFeatures.PromptSpecification promptSpec : this.prompts.values()) {
promptList.add(promptSpec.getPrompt());
}
return CompletableFuture.completedFuture(new McpSchema.ListPromptsResult(promptList, null));
};
}
private McpServerSession.RequestHandler<McpSchema.GetPromptResult> promptsGetRequestHandler() {
return (exchange, params) -> {
McpSchema.GetPromptRequest promptRequest = objectMapper.convertValue(params,
new TypeReference<McpSchema.GetPromptRequest>() {
});
McpServerFeatures.PromptSpecification specification = this.prompts.get(promptRequest.getName());
if (specification != null) {
return specification.getPromptHandler().apply(exchange, promptRequest);
}
CompletableFuture<McpSchema.GetPromptResult> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Prompt not found: " + promptRequest.getName()));
return future;
};
}
// ---------------------------------------
// Logging Management
// ---------------------------------------
public CompletableFuture<Void> loggingNotification(
McpSchema.LoggingMessageNotification loggingMessageNotification) {
Assert.notNull(loggingMessageNotification, "Logging message must not be null");
if (loggingMessageNotification.getLevel().level() < minLoggingLevel.level()) {
return CompletableFuture.completedFuture(null);
}
return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_MESSAGE,
loggingMessageNotification);
}
private McpServerSession.RequestHandler<Map<String, Object>> setLoggerRequestHandler() {
return (exchange, params) -> {
try {
McpSchema.SetLevelRequest request = this.objectMapper.convertValue(params,
McpSchema.SetLevelRequest.class);
this.minLoggingLevel = request.getLevel();
return CompletableFuture.completedFuture(Collections.emptyMap());
}
catch (Exception e) {
CompletableFuture<Map<String, Object>> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("An error occurred while processing a request to set the log level: " + e.getMessage()));
return future;
}
};
}
// ---------------------------------------
// Sampling
// ---------------------------------------
public void setProtocolVersions(List<String> protocolVersions) {
this.protocolVersions = protocolVersions;
}
}

View File

@ -0,0 +1,214 @@
package com.taobao.arthas.mcp.server.protocol.server;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.spec.McpError;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.LoggingLevel;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.LoggingMessageNotification;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerSession;
import com.taobao.arthas.mcp.server.util.Assert;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicLong;
/**
* Represents the interaction between MCP server and client. Provides methods for communication, logging, and context management.
*
* @author Yeaury
* <p>
* McpNettyServerExchange provides various methods for communicating with the client, including:
* <ul>
* <li>Sending requests and notifications
* <li>Getting client capabilities and information
* <li>Handling logging notifications
* <li>Creating client messages
* <li>Managing root directories
* </ul>
* <p>
* Each exchange object is associated with a specific client session, providing context and capabilities for that session.
*/
public class McpNettyServerExchange {
private static final Logger logger = LoggerFactory.getLogger(McpNettyServerExchange.class);
private final McpServerSession session;
private final McpSchema.ClientCapabilities clientCapabilities;
private final McpSchema.Implementation clientInfo;
private volatile LoggingLevel minLoggingLevel = LoggingLevel.INFO;
private static final TypeReference<McpSchema.CreateMessageResult> CREATE_MESSAGE_RESULT_TYPE_REF =
new TypeReference<McpSchema.CreateMessageResult>() {
};
private static final TypeReference<McpSchema.ListRootsResult> LIST_ROOTS_RESULT_TYPE_REF =
new TypeReference<McpSchema.ListRootsResult>() {
};
/**
* Create a new server exchange object.
* @param session Session associated with the client
* @param clientCapabilities Client capabilities
* @param clientInfo Client information
*/
public McpNettyServerExchange(McpServerSession session, McpSchema.ClientCapabilities clientCapabilities,
McpSchema.Implementation clientInfo) {
Assert.notNull(session, "Session cannot be null");
this.session = session;
this.clientCapabilities = clientCapabilities;
this.clientInfo = clientInfo;
logger.debug("Created new server exchange, session ID: {}, client: {}", session.getId(), clientInfo);
}
/**
* Create a new server exchange object. This constructor is used when the session exists but client capabilities and information are unknown.
* For example, for exchange objects created in request handlers.
* @param session Server session
* @param objectMapper JSON object mapper
*/
public McpNettyServerExchange(McpServerSession session, ObjectMapper objectMapper) {
Assert.notNull(session, "Session cannot be null");
this.session = session;
this.clientCapabilities = null;
this.clientInfo = null;
logger.debug("Created new server exchange, session ID: {}, client info unknown", session.getId());
}
/**
* Get client capabilities.
* @return Client capabilities
*/
public McpSchema.ClientCapabilities getClientCapabilities() {
return this.clientCapabilities;
}
/**
* Get client information.
* @return Client information
*/
public McpSchema.Implementation getClientInfo() {
return this.clientInfo;
}
/**
* Send a notification with parameters to the client.
* @param method Notification method name to send to the client
* @param params Parameters to send with the notification
* @return A CompletableFuture that completes when the notification is sent
*/
public CompletableFuture<Void> sendNotification(String method, Object params) {
Assert.hasText(method, "Method name cannot be empty");
logger.debug("Sending notification to client - method: {}, session: {}", method, this.session.getId());
return this.session.sendNotification(method, params).whenComplete((result, error) -> {
if (error != null) {
logger.error("Notification failed - method: {}, session: {}, error: {}", method, this.session.getId(), error.getMessage());
}
else {
logger.debug("Notification sent successfully - method: {}, session: {}", method, this.session.getId());
}
});
}
/**
* Create a new message using client sampling capability. MCP provides a standardized way for servers to request
* LLM sampling ("completion" or "generation") through the client. This flow allows clients to maintain control
* over model access, selection, and permissions while enabling servers to leverage AI capabilitieswithout server
* API keys. Servers can request text or image-based interactions and can optionally include context from the MCP
* server in their prompts.
* @param createMessageRequest Request to create a new message
* @return A CompletableFuture that completes when the message is created
*/
public CompletableFuture<McpSchema.CreateMessageResult> createMessage(
McpSchema.CreateMessageRequest createMessageRequest) {
if (this.clientCapabilities == null) {
logger.error("Client not initialized, cannot create message");
CompletableFuture<McpSchema.CreateMessageResult> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Client must be initialized first. Please call initialize method!"));
return future;
}
if (this.clientCapabilities.getSampling() == null) {
logger.error("Client not configured with sampling capability, cannot create message");
CompletableFuture<McpSchema.CreateMessageResult> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("Client must be configured with sampling capability"));
return future;
}
logger.debug("Creating client message, session ID: {}", this.session.getId());
return this.session
.sendRequest(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, createMessageRequest, CREATE_MESSAGE_RESULT_TYPE_REF)
.whenComplete((result, error) -> {
if (error != null) {
logger.error("Failed to create message, session ID: {}, error: {}", this.session.getId(), error.getMessage());
}
else {
logger.debug("Message created successfully, session ID: {}", this.session.getId());
}
});
}
/**
* Get a list of all root directories provided by the client.
* @return Sends out the CompletableFuture of the root list result
*/
public CompletableFuture<McpSchema.ListRootsResult> listRoots() {
return this.listRoots(null);
}
/**
* Get the client-provided list of pagination roots.
* Optional pagination cursor @param cursor for the previous list request
* @return Emits a CompletableFuture containing the results of the root list
*/
public CompletableFuture<McpSchema.ListRootsResult> listRoots(String cursor) {
logger.debug("Requesting root list, session ID: {}, cursor: {}", this.session.getId(), cursor);
return this.session
.sendRequest(McpSchema.METHOD_ROOTS_LIST, new McpSchema.PaginatedRequest(cursor),
LIST_ROOTS_RESULT_TYPE_REF)
.whenComplete((result, error) -> {
if (error != null) {
logger.error("Failed to get root list, session ID: {}, error: {}", this.session.getId(), error.getMessage());
}
else {
logger.debug("Root list retrieved successfully, session ID: {}", this.session.getId());
}
});
}
public CompletableFuture<Void> loggingNotification(LoggingMessageNotification loggingMessageNotification) {
if (loggingMessageNotification == null) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new McpError("日志消息不能为空"));
return future;
}
if (this.isNotificationForLevelAllowed(loggingMessageNotification.getLevel())) {
return this.session.sendNotification(McpSchema.METHOD_NOTIFICATION_MESSAGE, loggingMessageNotification)
.whenComplete((result, error) -> {
if (error != null) {
logger.error("Failed to send logging notification, level: {}, session ID: {}, error: {}", loggingMessageNotification.getLevel(),
this.session.getId(), error.getMessage());
}
});
}
return CompletableFuture.completedFuture(null);
}
public void setMinLoggingLevel(LoggingLevel minLoggingLevel) {
Assert.notNull(minLoggingLevel, "最低日志级别不能为空");
logger.debug("Setting minimum logging level: {}, session ID: {}", minLoggingLevel, this.session.getId());
this.minLoggingLevel = minLoggingLevel;
}
private boolean isNotificationForLevelAllowed(LoggingLevel loggingLevel) {
return loggingLevel.level() >= this.minLoggingLevel.level();
}
}

View File

@ -0,0 +1,242 @@
package com.taobao.arthas.mcp.server.protocol.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.spec.*;
import com.taobao.arthas.mcp.server.util.Assert;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.function.BiFunction;
/**
* MCP server interface and builder for Netty-based implementation.
*
* @author Yeaury
*/
public interface McpServer {
static NettySpecification netty(McpServerTransportProvider transportProvider) {
return new NettySpecification(transportProvider);
}
/**
* serverSpecification
*/
class NettySpecification {
private static final McpSchema.Implementation DEFAULT_SERVER_INFO = new McpSchema.Implementation("mcp-server",
"1.0.0");
private final McpServerTransportProvider transportProvider;
private ObjectMapper objectMapper;
private McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO;
private McpSchema.ServerCapabilities serverCapabilities;
private String instructions;
/**
* The Model Context Protocol (MCP) allows servers to expose tools that can be
* invoked by language models. Tools enable models to interact with external
* systems, such as querying databases, calling APIs, or performing computations.
* Each tool is uniquely identified by a name and includes metadata describing its
* schema.
*/
private final List<McpServerFeatures.ToolSpecification> tools = new ArrayList<>();
/**
* The Model Context Protocol (MCP) provides a standardized way for servers to
* expose resources to clients. Resources allow servers to share data that
* provides context to language models, such as files, database schemas, or
* application-specific information. Each resource is uniquely identified by a
* URI.
*/
private final Map<String, McpServerFeatures.ResourceSpecification> resources = new HashMap<>();
private final List<McpSchema.ResourceTemplate> resourceTemplates = new ArrayList<>();
/**
* The Model Context Protocol (MCP) provides a standardized way for servers to
* expose prompt templates to clients. Prompts allow servers to provide structured
* messages and instructions for interacting with language models. Clients can
* discover available prompts, retrieve their contents, and provide arguments to
* customize them.
*/
private final Map<String, McpServerFeatures.PromptSpecification> prompts = new HashMap<>();
private final List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> rootsChangeHandlers = new ArrayList<>();
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
private NettySpecification(McpServerTransportProvider transportProvider) {
Assert.notNull(transportProvider, "Transport provider must not be null");
this.transportProvider = transportProvider;
}
/**
* Sets the server implementation information that will be shared with clients
* during connection initialization. This helps with version compatibility,
* debugging, and server identification.
* @param serverInfo The server implementation details including name and version.
* Must not be null.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if serverInfo is null
*/
public NettySpecification serverInfo(McpSchema.Implementation serverInfo) {
Assert.notNull(serverInfo, "Server info must not be null");
this.serverInfo = serverInfo;
return this;
}
public NettySpecification requestTimeout(Duration requestTimeout) {
Assert.notNull(requestTimeout, "Request timeout must not be null");
this.requestTimeout = requestTimeout;
return this;
}
public NettySpecification serverInfo(String name, String version) {
Assert.hasText(name, "Name must not be null or empty");
Assert.hasText(version, "Version must not be null or empty");
this.serverInfo = new McpSchema.Implementation(name, version);
return this;
}
public NettySpecification instructions(String instructions) {
this.instructions = instructions;
return this;
}
public NettySpecification capabilities(McpSchema.ServerCapabilities serverCapabilities) {
this.serverCapabilities = serverCapabilities;
return this;
}
public NettySpecification tool(McpSchema.Tool tool,
BiFunction<McpNettyServerExchange, Map<String, Object>, CompletableFuture<McpSchema.CallToolResult>> handler) {
Assert.notNull(tool, "Tool must not be null");
Assert.notNull(handler, "Handler must not be null");
this.tools.add(new McpServerFeatures.ToolSpecification(tool, handler));
return this;
}
public NettySpecification tools(List<McpServerFeatures.ToolSpecification> toolRegistrations) {
Assert.notNull(toolRegistrations, "Tool handlers list must not be null");
this.tools.addAll(toolRegistrations);
return this;
}
public NettySpecification tools(McpServerFeatures.ToolSpecification... toolRegistrations) {
for (McpServerFeatures.ToolSpecification tool : toolRegistrations) {
this.tools.add(tool);
}
return this;
}
public NettySpecification resources(Map<String, McpServerFeatures.ResourceSpecification> resourceSpecifications) {
Assert.notNull(resourceSpecifications, "Resource handlers map must not be null");
this.resources.putAll(resourceSpecifications);
return this;
}
public NettySpecification resources(List<McpServerFeatures.ResourceSpecification> resourceSpecifications) {
Assert.notNull(resourceSpecifications, "Resource handlers list must not be null");
for (McpServerFeatures.ResourceSpecification resource : resourceSpecifications) {
this.resources.put(resource.getResource().getUri(), resource);
}
return this;
}
public NettySpecification resources(McpServerFeatures.ResourceSpecification... resourceSpecifications) {
Assert.notNull(resourceSpecifications, "Resource handlers list must not be null");
for (McpServerFeatures.ResourceSpecification resource : resourceSpecifications) {
this.resources.put(resource.getResource().getUri(), resource);
}
return this;
}
public NettySpecification resourceTemplates(List<McpSchema.ResourceTemplate> resourceTemplates) {
Assert.notNull(resourceTemplates, "Resource templates must not be null");
this.resourceTemplates.addAll(resourceTemplates);
return this;
}
public NettySpecification resourceTemplates(McpSchema.ResourceTemplate... resourceTemplates) {
Assert.notNull(resourceTemplates, "Resource templates must not be null");
this.resourceTemplates.addAll(Arrays.asList(resourceTemplates));
return this;
}
public NettySpecification prompts(Map<String, McpServerFeatures.PromptSpecification> prompts) {
Assert.notNull(prompts, "Prompts map must not be null");
this.prompts.putAll(prompts);
return this;
}
public NettySpecification prompts(List<McpServerFeatures.PromptSpecification> prompts) {
Assert.notNull(prompts, "Prompts map must not be null");
for (McpServerFeatures.PromptSpecification prompt : prompts) {
this.prompts.put(prompt.getPrompt().getName(), prompt);
}
return this;
}
public NettySpecification prompts(McpServerFeatures.PromptSpecification... prompts) {
Assert.notNull(prompts, "Prompts map must not be null");
for (McpServerFeatures.PromptSpecification prompt : prompts) {
this.prompts.put(prompt.getPrompt().getName(), prompt);
}
return this;
}
public NettySpecification rootsChangeHandler(
BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>> handler) {
Assert.notNull(handler, "Consumer must not be null");
this.rootsChangeHandlers.add(handler);
return this;
}
public NettySpecification rootsChangeHandlers(
List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> handlers) {
Assert.notNull(handlers, "Handlers list must not be null");
this.rootsChangeHandlers.addAll(handlers);
return this;
}
public NettySpecification rootsChangeHandlers(
@SuppressWarnings("unchecked") BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>... handlers) {
Assert.notNull(handlers, "Handlers list must not be null");
return this.rootsChangeHandlers(Arrays.asList(handlers));
}
public NettySpecification objectMapper(ObjectMapper objectMapper) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
this.objectMapper = objectMapper;
return this;
}
public McpNettyServer build() {
ObjectMapper mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
return new McpNettyServer(
this.transportProvider,
mapper,
this.requestTimeout,
new McpServerFeatures.McpServerConfig(
this.serverInfo,
this.serverCapabilities,
this.tools,
this.resources,
this.resourceTemplates,
this.prompts,
this.rootsChangeHandlers,
this.instructions
)
);
}
}
}

View File

@ -0,0 +1,219 @@
package com.taobao.arthas.mcp.server.protocol.server;
import com.taobao.arthas.mcp.server.protocol.spec.*;
import com.taobao.arthas.mcp.server.util.Assert;
import com.taobao.arthas.mcp.server.util.Utils;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.BiFunction;
/**
* MCP server function specification, the server can choose the supported features.
* This implementation only provides an asynchronous API.
*
* @author Yeaury
*/
public class McpServerFeatures {
public static class McpServerConfig {
private final McpSchema.Implementation serverInfo;
private final McpSchema.ServerCapabilities serverCapabilities;
private final List<ToolSpecification> tools;
private final Map<String, ResourceSpecification> resources;
private final List<McpSchema.ResourceTemplate> resourceTemplates;
private final Map<String, PromptSpecification> prompts;
private final List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> rootsChangeConsumers;
private final String instructions;
public McpServerConfig(
McpSchema.Implementation serverInfo,
McpSchema.ServerCapabilities serverCapabilities,
List<ToolSpecification> tools,
Map<String, ResourceSpecification> resources,
List<McpSchema.ResourceTemplate> resourceTemplates,
Map<String, PromptSpecification> prompts,
List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> rootsChangeConsumers,
String instructions) {
Assert.notNull(serverInfo, "The server information cannot be empty");
// If serverCapabilities is empty, the appropriate capability configuration
// is automatically built based on the provided capabilities
if (serverCapabilities == null) {
serverCapabilities = new McpSchema.ServerCapabilities(
null, // experimental
new McpSchema.ServerCapabilities.LoggingCapabilities(),
!Utils.isEmpty(prompts) ? new McpSchema.ServerCapabilities.PromptCapabilities(false) : null,
!Utils.isEmpty(resources) ? new McpSchema.ServerCapabilities.ResourceCapabilities(false, false) : null,
!Utils.isEmpty(tools) ? new McpSchema.ServerCapabilities.ToolCapabilities(false) : null);
}
this.tools = (tools != null) ? tools : Collections.emptyList();
this.resources = (resources != null) ? resources : Collections.emptyMap();
this.resourceTemplates = (resourceTemplates != null) ? resourceTemplates : Collections.emptyList();
this.prompts = (prompts != null) ? prompts : Collections.emptyMap();
this.rootsChangeConsumers = (rootsChangeConsumers != null) ? rootsChangeConsumers : Collections.emptyList();
this.serverInfo = serverInfo;
this.serverCapabilities = serverCapabilities;
this.instructions = instructions;
}
public McpSchema.Implementation getServerInfo() {
return serverInfo;
}
public McpSchema.ServerCapabilities getServerCapabilities() {
return serverCapabilities;
}
public List<ToolSpecification> getTools() {
return tools;
}
public Map<String, ResourceSpecification> getResources() {
return resources;
}
public List<McpSchema.ResourceTemplate> getResourceTemplates() {
return resourceTemplates;
}
public Map<String, PromptSpecification> getPrompts() {
return prompts;
}
public List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> getRootsChangeConsumers() {
return rootsChangeConsumers;
}
public String getInstructions() {
return instructions;
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private McpSchema.Implementation serverInfo;
private McpSchema.ServerCapabilities serverCapabilities;
private final List<ToolSpecification> tools = new ArrayList<>();
private final Map<String, ResourceSpecification> resources = new HashMap<>();
private final List<McpSchema.ResourceTemplate> resourceTemplates = new ArrayList<>();
private final Map<String, PromptSpecification> prompts = new HashMap<>();
private final List<BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>>> rootsChangeConsumers = new ArrayList<>();
private String instructions;
public Builder serverInfo(McpSchema.Implementation serverInfo) {
this.serverInfo = serverInfo;
return this;
}
public Builder serverCapabilities(McpSchema.ServerCapabilities serverCapabilities) {
this.serverCapabilities = serverCapabilities;
return this;
}
public Builder addTool(ToolSpecification tool) {
this.tools.add(tool);
return this;
}
public Builder addResource(String key, ResourceSpecification resource) {
this.resources.put(key, resource);
return this;
}
public Builder addResourceTemplate(McpSchema.ResourceTemplate template) {
this.resourceTemplates.add(template);
return this;
}
public Builder addPrompt(String key, PromptSpecification prompt) {
this.prompts.put(key, prompt);
return this;
}
public Builder addRootsChangeConsumer(
BiFunction<McpNettyServerExchange, List<McpSchema.Root>, CompletableFuture<Void>> consumer) {
this.rootsChangeConsumers.add(consumer);
return this;
}
public Builder instructions(String instructions) {
this.instructions = instructions;
return this;
}
public McpServerConfig build() {
return new McpServerConfig(serverInfo, serverCapabilities, tools, resources, resourceTemplates, prompts,
rootsChangeConsumers, instructions);
}
}
}
public static class ToolSpecification {
private final McpSchema.Tool tool;
private final BiFunction<McpNettyServerExchange, Map<String, Object>, CompletableFuture<McpSchema.CallToolResult>> call;
public ToolSpecification(
McpSchema.Tool tool,
BiFunction<McpNettyServerExchange, Map<String, Object>, CompletableFuture<McpSchema.CallToolResult>> call) {
this.tool = tool;
this.call = call;
}
public McpSchema.Tool getTool() {
return tool;
}
public BiFunction<McpNettyServerExchange, Map<String, Object>, CompletableFuture<McpSchema.CallToolResult>> getCall() {
return call;
}
}
public static class ResourceSpecification {
private final McpSchema.Resource resource;
private final BiFunction<McpNettyServerExchange, McpSchema.ReadResourceRequest, CompletableFuture<McpSchema.ReadResourceResult>> readHandler;
public ResourceSpecification(
McpSchema.Resource resource,
BiFunction<McpNettyServerExchange, McpSchema.ReadResourceRequest, CompletableFuture<McpSchema.ReadResourceResult>> readHandler) {
this.resource = resource;
this.readHandler = readHandler;
}
public McpSchema.Resource getResource() {
return resource;
}
public BiFunction<McpNettyServerExchange, McpSchema.ReadResourceRequest, CompletableFuture<McpSchema.ReadResourceResult>> getReadHandler() {
return readHandler;
}
}
public static class PromptSpecification {
private final McpSchema.Prompt prompt;
private final BiFunction<McpNettyServerExchange, McpSchema.GetPromptRequest, CompletableFuture<McpSchema.GetPromptResult>> promptHandler;
public PromptSpecification(
McpSchema.Prompt prompt,
BiFunction<McpNettyServerExchange, McpSchema.GetPromptRequest, CompletableFuture<McpSchema.GetPromptResult>> promptHandler) {
this.prompt = prompt;
this.promptHandler = promptHandler;
}
public McpSchema.Prompt getPrompt() {
return prompt;
}
public BiFunction<McpNettyServerExchange, McpSchema.GetPromptRequest, CompletableFuture<McpSchema.GetPromptResult>> getPromptHandler() {
return promptHandler;
}
}
}

View File

@ -0,0 +1,379 @@
package com.taobao.arthas.mcp.server.protocol.server.handler;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerSession;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerTransport;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.handler.codec.http.*;
import io.netty.util.CharsetUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
/**
* MCP (Model Context Protocol) Request Handler
* Used to handle model context protocol requests for Arthas.
*
* @author Yeaury
*/
public class McpRequestHandler {
private static final Logger logger = LoggerFactory.getLogger(McpRequestHandler.class);
/** Message event type */
public static final String MESSAGE_EVENT_TYPE = "message";
/** Endpoint event type */
public static final String ENDPOINT_EVENT_TYPE = "endpoint";
/** Active session mapping: session ID -> session object */
private final Map<String, McpServerSession> sessions = new ConcurrentHashMap<>();
/** SSE connection mapping: session ID -> Channel */
private final Map<String, Channel> sseChannels = new ConcurrentHashMap<>();
/** Message endpoint path */
private final String messageEndpoint;
/** SSE endpoint path */
private final String sseEndpoint;
/** Object mapper for JSON serialization/deserialization */
private final ObjectMapper objectMapper;
/** Session factory */
private McpServerSession.Factory sessionFactory;
/**
* Create a new McpRequestHandler instance.
* @param messageEndpoint Message endpoint path, e.g. "/mcp/message"
* @param sseEndpoint SSE endpoint path, e.g. "/mcp"
* @param objectMapper Object mapper for JSON serialization/deserialization
*/
public McpRequestHandler(String messageEndpoint,
String sseEndpoint, ObjectMapper objectMapper) {
this.messageEndpoint = messageEndpoint;
this.sseEndpoint = sseEndpoint;
this.objectMapper = objectMapper;
}
public void handle(ChannelHandlerContext ctx, FullHttpRequest request){
String uri = request.uri();
if (request.method() == HttpMethod.GET && uri.endsWith(sseEndpoint)) {
handleSseRequest(ctx);
return;
}
if (request.method() == HttpMethod.POST && uri.contains(messageEndpoint)) {
handleMessageRequest(ctx, request);
return;
}
if (request.method() == HttpMethod.OPTIONS) {
sendOptionsResponse(ctx);
return;
}
sendNotFoundResponse(ctx);
}
private void handleSseRequest(ChannelHandlerContext ctx) {
String sessionId = UUID.randomUUID().toString();
HttpResponse response = new DefaultHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK);
response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/event-stream");
response.headers().set(HttpHeaderNames.CACHE_CONTROL, "no-cache");
response.headers().set(HttpHeaderNames.CONNECTION, HttpHeaderValues.KEEP_ALIVE);
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*");
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_HEADERS, "*");
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_METHODS, "GET, POST, OPTIONS");
HttpUtil.setTransferEncodingChunked(response, true);
ctx.writeAndFlush(response);
sseChannels.put(sessionId, ctx.channel());
HttpServerTransport transport = new HttpServerTransport(sessionId, ctx.channel());
McpServerSession session = sessionFactory.create(transport);
sessions.put(sessionId, session);
String endpointInfo = messageEndpoint + "?sessionId=" + sessionId;
sendSseEvent(ctx, ENDPOINT_EVENT_TYPE, endpointInfo);
logger.debug("SSE connection established, session ID: {}", sessionId);
}
private void handleMessageRequest(ChannelHandlerContext ctx, FullHttpRequest request) {
String uri = request.uri();
String sessionIdParam = "sessionId=";
int sessionIdStart = uri.indexOf(sessionIdParam);
if (sessionIdStart == -1) {
sendErrorResponse(ctx, HttpResponseStatus.BAD_REQUEST, "Missing sessionId parameter");
return;
}
String sessionId = uri.substring(sessionIdStart + sessionIdParam.length());
McpServerSession session = sessions.get(sessionId);
if (session == null) {
sendErrorResponse(ctx, HttpResponseStatus.NOT_FOUND, "Session not found");
return;
}
ByteBuf content = request.content();
String jsonContent = content.toString(CharsetUtil.UTF_8);
try {
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, jsonContent);
session.handle(message).whenComplete((result, ex) -> {
if (ex != null) {
logger.error("Error processing message", ex);
sendErrorResponse(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, ex.getMessage());
}
else {
sendSuccessResponse(ctx);
}
});
}
catch (Exception e) {
logger.error("Error parsing message", e);
sendErrorResponse(ctx, HttpResponseStatus.BAD_REQUEST, "Invalid message format");
}
}
private void sendSseEvent(ChannelHandlerContext ctx, String eventType, String data) {
StringBuilder eventBuilder = new StringBuilder();
eventBuilder.append("event: ").append(eventType).append("\n");
eventBuilder.append("data: ").append(data).append("\n\n");
ByteBuf buffer = Unpooled.copiedBuffer(eventBuilder.toString(), CharsetUtil.UTF_8);
ctx.writeAndFlush(new DefaultHttpContent(buffer));
}
private void sendOptionsResponse(ChannelHandlerContext ctx) {
FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK);
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*");
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_METHODS, "GET, POST, OPTIONS");
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_HEADERS, "*");
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_MAX_AGE, "86400");
response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0);
ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
}
private void sendSuccessResponse(ChannelHandlerContext ctx) {
FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK);
response.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/json");
response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0);
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*");
ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
}
private void sendErrorResponse(ChannelHandlerContext ctx, HttpResponseStatus status, String message) {
ByteBuf content = Unpooled.copiedBuffer("{\"error\":\"" + message + "\"}", CharsetUtil.UTF_8);
FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, status, content);
response.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/json");
response.headers().set(HttpHeaderNames.CONTENT_LENGTH, content.readableBytes());
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*");
ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
}
private void sendNotFoundResponse(ChannelHandlerContext ctx) {
ByteBuf content = Unpooled.copiedBuffer("{\"error\":\"Resource not found\"}", CharsetUtil.UTF_8);
FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.NOT_FOUND,
content);
response.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/json");
response.headers().set(HttpHeaderNames.CONTENT_LENGTH, content.readableBytes());
response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*");
ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
}
public CompletableFuture<Void> notifyClients(String method, Object params) {
CompletableFuture<Void> future = new CompletableFuture<>();
if (sessions.isEmpty()) {
logger.debug("No active sessions to broadcast message");
future.complete(null);
return future;
}
logger.debug("Attempting to broadcast message to {} active sessions", sessions.size());
CompletableFuture<?>[] futures = sessions.values()
.stream()
.map(session -> session.sendNotification(method, params).exceptionally(e -> {
logger.error("Failed to send message to session {}: {}", session.getId(), e.getMessage());
return null;
}))
.toArray(CompletableFuture[]::new);
CompletableFuture.allOf(futures).whenComplete((result, ex) -> {
if (ex != null) {
logger.error("Error broadcasting message", ex);
future.completeExceptionally(ex);
}
else {
future.complete(null);
}
});
return future;
}
public CompletableFuture<Void> closeGracefully() {
CompletableFuture<Void> future = new CompletableFuture<>();
try {
this.shutdown();
future.complete(null);
}
catch (Exception e) {
future.completeExceptionally(e);
}
return future;
}
private void shutdown() {
for (McpServerSession session : this.sessions.values()) {
try {
session.close();
}
catch (Exception e) {
logger.warn("Error closing session: {}", e.getMessage());
}
}
logger.info("MCP HTTP server closed");
}
public void setSessionFactory(McpServerSession.Factory sessionFactory) {
this.sessionFactory = sessionFactory;
}
public String getMessageEndpoint() {
return messageEndpoint;
}
public String getSseEndpoint() {
return sseEndpoint;
}
private class HttpServerTransport implements McpServerTransport {
private final String sessionId;
private final Channel channel;
public HttpServerTransport(String sessionId, Channel channel) {
this.sessionId = sessionId;
this.channel = channel;
}
public String getSessionId() {
return sessionId;
}
@Override
public Channel getChannel() {
return channel;
}
@Override
public CompletableFuture<Void> sendMessage(McpSchema.JSONRPCMessage message) {
CompletableFuture<Void> future = new CompletableFuture<>();
if (!channel.isActive()) {
future.completeExceptionally(new RuntimeException("Channel is not active"));
return future;
}
try {
String jsonMessage = objectMapper.writeValueAsString(message);
logger.debug("sendMessage: {}", jsonMessage);
StringBuilder eventBuilder = new StringBuilder();
eventBuilder.append("event: ").append(MESSAGE_EVENT_TYPE).append("\n");
eventBuilder.append("data: ").append(jsonMessage).append("\n\n");
ByteBuf buffer = Unpooled.copiedBuffer(eventBuilder.toString(), CharsetUtil.UTF_8);
channel.writeAndFlush(new DefaultHttpContent(buffer))
.addListener((ChannelFutureListener) channelFuture -> {
if (channelFuture.isSuccess()) {
future.complete(null);
}
else {
future.completeExceptionally(channelFuture.cause());
}
});
}
catch (Exception e) {
logger.error("Failed to send message to session {}: {}", sessionId, e.getMessage());
sessions.remove(sessionId);
future.completeExceptionally(e);
}
return future;
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return objectMapper.convertValue(data, typeRef);
}
@Override
public CompletableFuture<Void> closeGracefully() {
CompletableFuture<Void> future = new CompletableFuture<>();
try {
sessions.remove(sessionId);
sseChannels.remove(sessionId);
if (channel.isActive()) {
channel.close().addListener((ChannelFutureListener) channelFuture -> {
if (channelFuture.isSuccess()) {
future.complete(null);
}
else {
future.completeExceptionally(channelFuture.cause());
}
});
}
else {
future.complete(null);
}
}
catch (Exception e) {
future.completeExceptionally(e);
}
return future;
}
@Override
public void close() {
sessions.remove(sessionId);
sseChannels.remove(sessionId);
if (channel.isActive()) {
channel.close();
}
}
}
}

View File

@ -0,0 +1,90 @@
package com.taobao.arthas.mcp.server.protocol.server.transport;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerSession;
import com.taobao.arthas.mcp.server.protocol.spec.McpServerTransportProvider;
import com.taobao.arthas.mcp.server.util.Assert;
import java.util.concurrent.CompletableFuture;
/**
* HTTP MCP server transport provider implementation based on Netty,
* supporting server-to-client communication via SSE.
*
* @author Yeaury
*/
public class HttpNettyServerTransportProvider implements McpServerTransportProvider {
public static final String DEFAULT_SSE_ENDPOINT = "/mcp";
public static final String DEFAULT_MESSAGE_ENDPOINT = "/mcp/message";
private final McpRequestHandler mcpRequestHandler;
/**
* Create a new HttpServerTransportProvider instance.
* @param messageEndpoint Message endpoint path, e.g. "/mcp/message"
* @param sseEndpoint SSE endpoint path, e.g. "/mcp"
* @param objectMapper Object mapper for JSON serialization/deserialization
*/
public HttpNettyServerTransportProvider(String messageEndpoint, String sseEndpoint,
ObjectMapper objectMapper) {
Assert.hasText(messageEndpoint, "Message endpoint path cannot be empty");
Assert.hasText(sseEndpoint, "SSE endpoint path cannot be empty");
Assert.notNull(objectMapper, "Object mapper cannot be null");
this.mcpRequestHandler = new McpRequestHandler(messageEndpoint, sseEndpoint, objectMapper);
}
@Override
public void setSessionFactory(McpServerSession.Factory sessionFactory) {
this.mcpRequestHandler.setSessionFactory(sessionFactory);
}
@Override
public CompletableFuture<Void> notifyClients(String method, Object params) {
return mcpRequestHandler.notifyClients(method, params);
}
@Override
public CompletableFuture<Void> closeGracefully() {
return mcpRequestHandler.closeGracefully();
}
@Override
public McpRequestHandler getMcpRequestHandler() {
return mcpRequestHandler;
}
public static class Builder {
private String messageEndpoint = DEFAULT_MESSAGE_ENDPOINT;
private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
private ObjectMapper objectMapper = new ObjectMapper();
public Builder messageEndpoint(String messageEndpoint) {
this.messageEndpoint = messageEndpoint;
return this;
}
public Builder sseEndpoint(String sseEndpoint) {
this.sseEndpoint = sseEndpoint;
return this;
}
public Builder objectMapper(ObjectMapper objectMapper) {
this.objectMapper = objectMapper;
return this;
}
public HttpNettyServerTransportProvider build() {
return new HttpNettyServerTransportProvider(messageEndpoint, sseEndpoint, objectMapper);
}
}
public static Builder builder() {
return new Builder();
}
}

View File

@ -0,0 +1,27 @@
package com.taobao.arthas.mcp.server.protocol.spec;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.JSONRPCResponse.JSONRPCError;
/**
* Exception class for representing JSON-RPC errors in MCP protocol.
*
* @author Yeaury
*/
public class McpError extends RuntimeException {
private JSONRPCError jsonRpcError;
public McpError(JSONRPCError jsonRpcError) {
super(jsonRpcError.getMessage());
this.jsonRpcError = jsonRpcError;
}
public McpError(Object error) {
super(error.toString());
}
public JSONRPCError getJsonRpcError() {
return jsonRpcError;
}
}

View File

@ -0,0 +1,433 @@
package com.taobao.arthas.mcp.server.protocol.spec;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import com.fasterxml.jackson.core.type.TypeReference;
import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange;
import io.netty.channel.Channel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
public class McpServerSession implements McpSession {
private static final Logger logger = LoggerFactory.getLogger(McpServerSession.class);
private final ConcurrentHashMap<Object, CompletableFuture<McpSchema.JSONRPCResponse>> pendingResponses = new ConcurrentHashMap<>();
private final String id;
private final Duration requestTimeout;
private final AtomicLong requestCounter = new AtomicLong(0);
private final InitRequestHandler initRequestHandler;
private final InitNotificationHandler initNotificationHandler;
private final Map<String, RequestHandler<?>> requestHandlers;
private final Map<String, NotificationHandler> notificationHandlers;
private final McpServerTransport transport;
private final CompletableFuture<McpNettyServerExchange> exchangeFuture = new CompletableFuture<>();
private final AtomicReference<McpSchema.ClientCapabilities> clientCapabilities = new AtomicReference<>();
private final AtomicReference<McpSchema.Implementation> clientInfo = new AtomicReference<>();
private final Channel channel;
private static final int STATE_UNINITIALIZED = 0;
private static final int STATE_INITIALIZING = 1;
private static final int STATE_INITIALIZED = 2;
private final AtomicInteger state = new AtomicInteger(STATE_UNINITIALIZED);
/**
* Create a new server session.
* @param id Session ID
* @param requestTimeout request timeout
* @param transport layer used
* @param initHandler handlers that handle initialization requests
* @param initNotificationHandler handles the handler that initializes the notification
* @param requestHandlers request processor mappings
* @param notificationHandlers notification handler mapping
* @param channel Netty's Channel object
*/
public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
InitRequestHandler initHandler, InitNotificationHandler initNotificationHandler,
Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers,
Channel channel) {
this.id = id;
this.requestTimeout = requestTimeout;
this.transport = transport;
this.initRequestHandler = initHandler;
this.initNotificationHandler = initNotificationHandler;
this.requestHandlers = requestHandlers;
this.notificationHandlers = notificationHandlers;
this.channel = channel;
}
public String getId() {
return this.id;
}
/**
* Called after successfully initializing the sequence between the client and server,
* including client capabilities and information.
* @param clientCapabilities The capabilities provided by the connected client
* @param clientInfo Information about the connected client
*/
public void init(McpSchema.ClientCapabilities clientCapabilities, McpSchema.Implementation clientInfo) {
this.clientCapabilities.lazySet(clientCapabilities);
this.clientInfo.lazySet(clientInfo);
}
private String generateRequestId() {
return this.id + "-" + this.requestCounter.getAndIncrement();
}
@Override
public <T> CompletableFuture<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
String requestId = this.generateRequestId();
CompletableFuture<T> result = new CompletableFuture<>();
CompletableFuture<McpSchema.JSONRPCResponse> responseFuture = new CompletableFuture<>();
this.pendingResponses.put(requestId, responseFuture);
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
requestId, requestParams);
// 使用Netty的Channel进行异步发送
channel.eventLoop().execute(() -> {
transport.sendMessage(jsonrpcRequest).exceptionally(error -> {
pendingResponses.remove(requestId);
responseFuture.completeExceptionally(error);
return null;
});
});
// 设置超时处理
channel.eventLoop().schedule(() -> {
if (!responseFuture.isDone()) {
pendingResponses.remove(requestId);
responseFuture.completeExceptionally(
new RuntimeException("Request timed out after " + requestTimeout.toMillis() + "ms"));
}
}, requestTimeout.toMillis(), java.util.concurrent.TimeUnit.MILLISECONDS);
// 处理响应结果
responseFuture.thenAccept(jsonRpcResponse -> {
if (jsonRpcResponse.getError() != null) {
result.completeExceptionally(new McpError(jsonRpcResponse.getError()));
}
else {
if (typeRef.getType().equals(Void.class)) {
result.complete(null);
}
else {
result.complete(this.transport.unmarshalFrom(jsonRpcResponse.getResult(), typeRef));
}
}
}).exceptionally(error -> {
result.completeExceptionally(error);
return null;
});
return result;
}
@Override
public CompletableFuture<Void> sendNotification(String method, Object params) {
McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
method, params);
return CompletableFuture.supplyAsync(() -> {
channel.eventLoop().execute(() -> {
transport.sendMessage(jsonrpcNotification);
});
return null;
}, runnable -> channel.eventLoop().execute(runnable));
}
public CompletableFuture<Void> handle(McpSchema.JSONRPCMessage message) {
CompletableFuture<Void> result = new CompletableFuture<>();
try {
if (message instanceof McpSchema.JSONRPCResponse) {
handleResponse((McpSchema.JSONRPCResponse) message, result);
}
else if (message instanceof McpSchema.JSONRPCRequest) {
handleRequest((McpSchema.JSONRPCRequest) message, result);
}
else if (message instanceof McpSchema.JSONRPCNotification) {
handleNotification((McpSchema.JSONRPCNotification) message, result);
}
else {
logger.warn("Received unknown message type: {}", message);
result.complete(null);
}
}
catch (Exception e) {
logger.error("Error processing message", e);
result.completeExceptionally(e);
}
return result;
}
private void handleResponse(McpSchema.JSONRPCResponse response, CompletableFuture<Void> result) {
logger.debug("Received response: {}", response);
CompletableFuture<McpSchema.JSONRPCResponse> sink = pendingResponses.remove(response.getId());
if (sink == null) {
logger.warn("Received unexpected response with unknown ID: {}", response.getId());
}
else {
sink.complete(response);
}
result.complete(null);
}
private void handleRequest(McpSchema.JSONRPCRequest request, CompletableFuture<Void> result) {
logger.debug("Received request: {}", request);
handleIncomingRequest(request)
.thenCompose(this.transport::sendMessage)
.thenAccept(v -> result.complete(null))
.exceptionally(error -> {
McpSchema.JSONRPCResponse errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
error.getMessage(), null));
channel.eventLoop().execute(() -> {
this.transport.sendMessage(errorResponse)
.thenRun(() -> result.complete(null))
.exceptionally(e -> {
result.completeExceptionally(e);
return null;
});
});
return null;
});
}
private void handleNotification(McpSchema.JSONRPCNotification notification, CompletableFuture<Void> result) {
logger.debug("Received notification: {}", notification);
handleIncomingNotification(notification)
.thenAccept(v -> result.complete(null))
.exceptionally(error -> {
logger.error("Error processing notification: {}", error.getMessage());
result.complete(null);
return null;
});
}
private CompletableFuture<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCRequest request) {
String method = request.getMethod();
Object params = request.getParams();
if (McpSchema.METHOD_INITIALIZE.equals(method)) {
if (this.state.compareAndSet(STATE_UNINITIALIZED, STATE_INITIALIZING)) {
try {
McpSchema.InitializeRequest initRequest = this.transport.unmarshalFrom(params,
new TypeReference<McpSchema.InitializeRequest>() {
});
return this.initRequestHandler.handle(initRequest).thenApply(result -> {
this.state.set(STATE_INITIALIZED);
this.init(initRequest.getCapabilities(), initRequest.getClientInfo());
if (!this.exchangeFuture.isDone()) {
McpNettyServerExchange exchange = new McpNettyServerExchange(this,
initRequest.getCapabilities(), initRequest.getClientInfo());
this.exchangeFuture.complete(exchange);
}
return new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), result, null);
}).exceptionally(error -> {
this.state.set(STATE_UNINITIALIZED);
return new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
error.getMessage(), null));
});
}
catch (Exception e) {
this.state.set(STATE_UNINITIALIZED);
CompletableFuture<McpSchema.JSONRPCResponse> future = new CompletableFuture<>();
future.complete(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INVALID_PARAMS,
"invalid initialization parameter: " + e.getMessage(), null)));
return future;
}
}
else {
CompletableFuture<McpSchema.JSONRPCResponse> future = new CompletableFuture<>();
future.complete(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INVALID_REQUEST, "the session has been initialized",
null)));
return future;
}
}
// check if the session is initialized
if (this.state.get() != STATE_INITIALIZED) {
CompletableFuture<McpSchema.JSONRPCResponse> future = new CompletableFuture<>();
future.complete(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INVALID_REQUEST, "the session has not been initialized", null)));
return future;
}
// handle regular requests
RequestHandler<?> handler = this.requestHandlers.get(method);
if (handler == null) {
MethodNotFoundError error = getMethodNotFoundError(method);
CompletableFuture<McpSchema.JSONRPCResponse> future = new CompletableFuture<>();
future.complete(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND, error.getMessage(),
error.getData())));
return future;
}
// get the swap object and process the request
return this.exchangeFuture.thenCompose(exchange -> {
try {
@SuppressWarnings("unchecked")
RequestHandler<Object> typedHandler = (RequestHandler<Object>) handler;
return typedHandler.handle(exchange, params)
.thenApply(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), result,
null))
.exceptionally(error -> {
Throwable cause = error.getCause();
if (cause instanceof McpError) {
McpError mcpError = (McpError) cause;
return new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
mcpError.getJsonRpcError());
}
else {
return new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
error.getMessage(), null));
}
});
}
catch (Exception e) {
CompletableFuture<McpSchema.JSONRPCResponse> future = new CompletableFuture<>();
future.complete(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, e.getMessage(),
null)));
return future;
}
});
}
/**
* Handle incoming JSON-RPC notifications, routing them to the appropriate handler.
* @param notification incoming JSON-RPC notification
* @return a CompletableFuture that completes when the notification processing is finished.
* The CompletableFuture completes normally if the processing succeeds, or exceptionally
* if an error occurs during processing.
*/
private CompletableFuture<Void> handleIncomingNotification(McpSchema.JSONRPCNotification notification) {
String method = notification.getMethod();
Object params = notification.getParams();
if (McpSchema.METHOD_NOTIFICATION_INITIALIZED.equals(method)) {
return this.initNotificationHandler.handle();
}
if (this.state.get() != STATE_INITIALIZED) {
CompletableFuture<Void> future = new CompletableFuture<>();
future.completeExceptionally(new IllegalStateException("the session has not been initialized"));
return future;
}
NotificationHandler handler = this.notificationHandlers.get(method);
if (handler == null) {
logger.warn("No handler found for method: {}", method);
CompletableFuture<Void> future = new CompletableFuture<>();
future.complete(null);
return future;
}
return this.exchangeFuture.thenCompose(exchange -> handler.handle(exchange, params));
}
public static class MethodNotFoundError {
private final String method;
private final String message;
private final Object data;
public MethodNotFoundError(String method, String message, Object data) {
this.method = method;
this.message = message;
this.data = data;
}
public String getMethod() {
return method;
}
public String getMessage() {
return message;
}
public Object getData() {
return data;
}
}
static MethodNotFoundError getMethodNotFoundError(String method) {
Map<String, String> data = new HashMap<>();
data.put("method", method);
return new MethodNotFoundError(method, "method not found: " + method, data);
}
@Override
public CompletableFuture<Void> closeGracefully() {
return this.transport.closeGracefully();
}
@Override
public void close() {
this.transport.close();
}
public interface InitRequestHandler {
CompletableFuture<McpSchema.InitializeResult> handle(McpSchema.InitializeRequest initializeRequest);
}
public interface InitNotificationHandler {
CompletableFuture<Void> handle();
}
public interface NotificationHandler {
CompletableFuture<Void> handle(McpNettyServerExchange exchange, Object params);
}
public interface RequestHandler<T> {
CompletableFuture<T> handle(McpNettyServerExchange exchange, Object params);
}
@FunctionalInterface
public interface Factory {
McpServerSession create(McpServerTransport sessionTransport);
}
}

View File

@ -0,0 +1,13 @@
package com.taobao.arthas.mcp.server.protocol.spec;
import io.netty.channel.Channel;
/**
* Extends McpTransport to provide access to the underlying Netty Channel for server-side transport.
*
* @author Yeaury
*/
public interface McpServerTransport extends McpTransport {
Channel getChannel();
}

View File

@ -0,0 +1,26 @@
package com.taobao.arthas.mcp.server.protocol.spec;
import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler;
import java.util.concurrent.CompletableFuture;
/**
* Provides the abstraction for server-side transport providers in MCP protocol.
* Defines methods for session factory setup, client notification, and graceful shutdown.
*
* @author Yeaury
*/
public interface McpServerTransportProvider {
void setSessionFactory(McpServerSession.Factory sessionFactory);
CompletableFuture<Void> notifyClients(String method, Object params);
CompletableFuture<Void> closeGracefully();
default void close() {
closeGracefully();
}
McpRequestHandler getMcpRequestHandler();
}

View File

@ -0,0 +1,26 @@
package com.taobao.arthas.mcp.server.protocol.spec;
import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
/**
* Represents a server-side MCP session that manages bidirectional JSON-RPC communication with the client.
*
* @author Yeaury
*/
public interface McpSession {
<T> CompletableFuture<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef);
default CompletableFuture<Void> sendNotification(String method) {
return sendNotification(method, null);
}
CompletableFuture<Void> sendNotification(String method, Object params);
CompletableFuture<Void> closeGracefully();
void close();
}

View File

@ -0,0 +1,24 @@
package com.taobao.arthas.mcp.server.protocol.spec;
import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
/**
* Defines the transport abstraction for sending messages and unmarshalling data in MCP protocol.
*
* @author Yeaury
*/
public interface McpTransport {
CompletableFuture<Void> closeGracefully();
default void close() {
this.closeGracefully();
}
CompletableFuture<Void> sendMessage(McpSchema.JSONRPCMessage message);
<T> T unmarshalFrom(Object data, TypeReference<T> typeRef);
}

View File

@ -0,0 +1,190 @@
package com.taobao.arthas.mcp.server.tool;
import com.fasterxml.jackson.core.type.TypeReference;
import com.taobao.arthas.mcp.server.tool.definition.ToolDefinition;
import com.taobao.arthas.mcp.server.tool.execution.ToolCallResultConverter;
import com.taobao.arthas.mcp.server.tool.execution.ToolExecutionException;
import com.taobao.arthas.mcp.server.util.Assert;
import com.taobao.arthas.mcp.server.util.JsonParser;
import com.taobao.arthas.mcp.server.util.Utils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.Type;
import java.util.Arrays;
import java.util.Map;
import java.util.stream.Stream;
public class DefaultToolCallback implements ToolCallback {
private static final Logger logger = LoggerFactory.getLogger(DefaultToolCallback.class);
private final ToolDefinition toolDefinition;
private final Method toolMethod;
private final Object toolObject;
private final ToolCallResultConverter toolCallResultConverter;
public DefaultToolCallback(ToolDefinition toolDefinition, Method toolMethod,
Object toolObject, ToolCallResultConverter toolCallResultConverter) {
Assert.notNull(toolDefinition, "toolDefinition cannot be null");
Assert.notNull(toolMethod, "toolMethod cannot be null");
Assert.isTrue(Modifier.isStatic(toolMethod.getModifiers()) || toolObject != null,
"toolObject cannot be null for non-static methods");
this.toolDefinition = toolDefinition;
this.toolMethod = toolMethod;
this.toolObject = toolObject;
this.toolCallResultConverter = toolCallResultConverter;
}
@Override
public ToolDefinition getToolDefinition() {
return this.toolDefinition;
}
@Override
public String call(String toolInput) {
return call(toolInput, null);
}
@Override
public String call(String toolInput, ToolContext toolContext) {
Assert.hasText(toolInput, "toolInput cannot be null or empty");
logger.debug("Starting execution of tool: {}", this.toolDefinition.getName());
validateToolContextSupport(toolContext);
Map<String, Object> toolArguments = extractToolArguments(toolInput);
Object[] methodArguments = buildMethodArguments(toolArguments, toolContext);
Object result = callMethod(methodArguments);
logger.debug("Successful execution of tool: {}", this.toolDefinition.getName());
Type returnType = this.toolMethod.getGenericReturnType();
return this.toolCallResultConverter.convert(result, returnType);
}
private void validateToolContextSupport(ToolContext toolContext) {
boolean isNonEmptyToolContextProvided = toolContext != null && !Utils.isEmpty(toolContext.getContext());
boolean isToolContextAcceptedByMethod = Arrays.stream(this.toolMethod.getParameterTypes())
.anyMatch(type -> Utils.isAssignable(type, ToolContext.class));
if (isToolContextAcceptedByMethod && !isNonEmptyToolContextProvided) {
throw new IllegalArgumentException("ToolContext is required by the method as an argument");
}
}
private Map<String, Object> extractToolArguments(String toolInput) {
return JsonParser.fromJson(toolInput, new TypeReference<Map<String, Object>>() {
});
}
private Object[] buildMethodArguments(Map<String, Object> toolInputArguments, ToolContext toolContext) {
return Stream.of(this.toolMethod.getParameters()).map(parameter -> {
if (parameter.getType().isAssignableFrom(ToolContext.class)) {
return toolContext;
}
Object rawArgument = toolInputArguments.get(parameter.getName());
return buildTypedArgument(rawArgument, parameter.getParameterizedType());
}).toArray();
}
private Object buildTypedArgument(Object value, Type type) {
if (value == null) {
return null;
}
if (type instanceof Class<?>) {
return JsonParser.toTypedObject(value, (Class<?>) type);
}
String json = JsonParser.toJson(value);
return JsonParser.fromJson(json, type);
}
private Object callMethod(Object[] methodArguments) {
if (isObjectNotPublic() || isMethodNotPublic()) {
this.toolMethod.setAccessible(true);
}
Object result;
try {
result = this.toolMethod.invoke(this.toolObject, methodArguments);
}
catch (IllegalAccessException ex) {
throw new IllegalStateException("Could not access method: " + ex.getMessage(), ex);
}
catch (InvocationTargetException ex) {
throw new ToolExecutionException(this.toolDefinition, ex.getCause());
}
return result;
}
private boolean isObjectNotPublic() {
return this.toolObject != null && !Modifier.isPublic(this.toolObject.getClass().getModifiers());
}
private boolean isMethodNotPublic() {
return !Modifier.isPublic(this.toolMethod.getModifiers());
}
@Override
public String toString() {
return "MethodToolCallback{" + "toolDefinition=" + this.toolDefinition + '}';
}
public static Builder builder() {
return new Builder();
}
public static final class Builder {
private ToolDefinition toolDefinition;
private Method toolMethod;
private Object toolObject;
private ToolCallResultConverter toolCallResultConverter;
private Builder() {
}
public Builder toolDefinition(ToolDefinition toolDefinition) {
this.toolDefinition = toolDefinition;
return this;
}
public Builder toolMethod(Method toolMethod) {
this.toolMethod = toolMethod;
return this;
}
public Builder toolObject(Object toolObject) {
this.toolObject = toolObject;
return this;
}
public Builder toolCallResultConverter(ToolCallResultConverter toolCallResultConverter) {
this.toolCallResultConverter = toolCallResultConverter;
return this;
}
public DefaultToolCallback build() {
return new DefaultToolCallback(this.toolDefinition, this.toolMethod, this.toolObject, this.toolCallResultConverter);
}
}
}

View File

@ -0,0 +1,170 @@
package com.taobao.arthas.mcp.server.tool;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.definition.ToolDefinition;
import com.taobao.arthas.mcp.server.tool.definition.ToolDefinitions;
import com.taobao.arthas.mcp.server.tool.execution.DefaultToolCallResultConverter;
import com.taobao.arthas.mcp.server.tool.execution.ToolCallResultConverter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.net.JarURLConnection;
import java.net.URL;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
/**
* Default tool callback provider implementation
* <p>
* Scan methods with @Tool annotations in the classpath and register them as tool callbacks
*/
public class DefaultToolCallbackProvider implements ToolCallbackProvider {
private static final Logger logger = LoggerFactory.getLogger(DefaultToolCallbackProvider.class);
private static final String DEFAULT_TOOL_BASE_PACKAGE = "com.taobao.arthas.mcp.server.tool.function";
private final ToolCallResultConverter toolCallResultConverter;
private ToolCallback[] toolCallbacks;
private String toolBasePackage;
public DefaultToolCallbackProvider() {
this.toolCallResultConverter = new DefaultToolCallResultConverter();
this.toolBasePackage = DEFAULT_TOOL_BASE_PACKAGE;
}
public void setToolBasePackage(String toolBasePackage) {
this.toolBasePackage = toolBasePackage;
}
@Override
public ToolCallback[] getToolCallbacks() {
if (toolCallbacks == null) {
synchronized (this) {
if (toolCallbacks == null) {
toolCallbacks = scanForToolCallbacks();
}
}
}
return toolCallbacks;
}
private ToolCallback[] scanForToolCallbacks() {
List<ToolCallback> callbacks = new ArrayList<>();
try {
logger.info("Starting to scan for tool callbacks in package: {}", toolBasePackage);
scanPackageForToolMethods(toolBasePackage, callbacks);
logger.info("Found {} tool callbacks", callbacks.size());
} catch (Exception e) {
logger.error("Failed to scan for tool callbacks: {}", e.getMessage(), e);
}
return callbacks.toArray(new ToolCallback[0]);
}
private void scanPackageForToolMethods(String packageName, List<ToolCallback> callbacks) throws IOException {
String packageDirName = packageName.replace('.', '/');
ClassLoader classLoader = DefaultToolCallbackProvider.class.getClassLoader();
logger.info("Using classloader: {} for scanning package: {}", classLoader, packageName);
Enumeration<URL> resources = classLoader.getResources(packageDirName);
if (!resources.hasMoreElements()) {
logger.warn("No resources found for package: {}", packageName);
return;
}
while (resources.hasMoreElements()) {
URL resource = resources.nextElement();
String protocol = resource.getProtocol();
logger.info("Found resource: {} with protocol: {}", resource, protocol);
if ("file".equals(protocol)) {
String filePath = URLDecoder.decode(resource.getFile(), StandardCharsets.UTF_8.name());
logger.info("Scanning directory: {}", filePath);
scanDirectory(new File(filePath), packageName, callbacks);
} else if ("jar".equals(protocol)) {
JarURLConnection jarConn = (JarURLConnection) resource.openConnection();
try (JarFile jarFile = jarConn.getJarFile()) {
logger.info("Scanning jar file: {}", jarFile.getName());
scanJarEntries(jarFile, packageDirName, callbacks);
}
} else {
logger.warn("Unsupported protocol: {} for resource: {}", protocol, resource);
}
}
}
private void scanDirectory(File directory, String packageName, List<ToolCallback> callbacks) {
if (!directory.exists() || !directory.isDirectory()) {
logger.warn("Directory does not exist or is not a directory: {}", directory);
return;
}
File[] files = directory.listFiles();
if (files == null) {
logger.warn("Failed to list files in directory: {}", directory);
return;
}
for (File file : files) {
if (file.isDirectory()) {
scanDirectory(file, packageName + "." + file.getName(), callbacks);
} else if (file.getName().endsWith(".class")) {
String className = packageName + "." + file.getName().substring(0, file.getName().length() - 6);
logger.debug("Processing class: {}", className);
processClass(className, callbacks);
}
}
}
private void scanJarEntries(JarFile jarFile, String packageDirName, List<ToolCallback> callbacks) {
Enumeration<JarEntry> entries = jarFile.entries();
while (entries.hasMoreElements()) {
JarEntry entry = entries.nextElement();
String name = entry.getName();
if (name.startsWith(packageDirName) && name.endsWith(".class")) {
String className = name.substring(0, name.length() - 6).replace('/', '.');
logger.debug("Processing jar entry: {}", className);
processClass(className, callbacks);
}
}
}
private void processClass(String className, List<ToolCallback> callbacks) {
try {
Class<?> clazz = Class.forName(className, false, DefaultToolCallbackProvider.class.getClassLoader());
if (clazz.isInterface() || clazz.isEnum() || clazz.isAnnotation()) {
return;
}
for (Method method : clazz.getDeclaredMethods()) {
if (method.isAnnotationPresent(Tool.class)) {
registerToolMethod(clazz, method, callbacks);
}
}
} catch (Throwable t) {
logger.warn("Error loading class {}: {}", className, t.getMessage(), t);
}
}
private void registerToolMethod(Class<?> clazz, Method method, List<ToolCallback> callbacks) {
try {
ToolDefinition toolDefinition = ToolDefinitions.from(method);
Object toolObject = Modifier.isStatic(method.getModifiers()) ? null : clazz.getDeclaredConstructor().newInstance();
ToolCallback callback = DefaultToolCallback.builder()
.toolDefinition(toolDefinition)
.toolMethod(method)
.toolObject(toolObject)
.toolCallResultConverter(toolCallResultConverter)
.build();
callbacks.add(callback);
logger.info("Registered tool: {} from class: {}", toolDefinition.getName(), clazz.getName());
} catch (Exception e) {
logger.error("Failed to register tool {}.{}, error: {}",
clazz.getName(), method.getName(), e.getMessage(), e);
}
}
}

View File

@ -0,0 +1,15 @@
package com.taobao.arthas.mcp.server.tool;
import com.taobao.arthas.mcp.server.tool.definition.ToolDefinition;
/**
* Define the basic behavior of the tool
*/
public interface ToolCallback {
ToolDefinition getToolDefinition();
String call(String toolInput);
String call(String toolInput, ToolContext toolContext);
}

View File

@ -0,0 +1,7 @@
package com.taobao.arthas.mcp.server.tool;
public interface ToolCallbackProvider {
ToolCallback[] getToolCallbacks();
}

View File

@ -0,0 +1,18 @@
package com.taobao.arthas.mcp.server.tool;
import java.util.Collections;
import java.util.Map;
public final class ToolContext {
private final Map<String, Object> context;
public ToolContext(Map<String, Object> context) {
this.context = Collections.unmodifiableMap(context);
}
public Map<String, Object> getContext() {
return this.context;
}
}

View File

@ -0,0 +1,16 @@
package com.taobao.arthas.mcp.server.tool.annotation;
import java.lang.annotation.*;
@Target({ ElementType.METHOD, ElementType.ANNOTATION_TYPE })
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface Tool {
String name() default "";
String description() default "";
boolean streamable() default false;
}

View File

@ -0,0 +1,20 @@
package com.taobao.arthas.mcp.server.tool.annotation;
import java.lang.annotation.*;
@Target({ ElementType.PARAMETER, ElementType.FIELD, ElementType.ANNOTATION_TYPE })
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface ToolParam {
/**
* Whether the tool argument is required.
*/
boolean required() default true;
/**
* The description of the tool argument.
*/
String description() default "";
}

View File

@ -0,0 +1,76 @@
package com.taobao.arthas.mcp.server.tool.definition;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
public class ToolDefinition {
private String name;
private String description;
private McpSchema.JsonSchema inputSchema;
private boolean streamable;
public ToolDefinition(String name, String description,
McpSchema.JsonSchema inputSchema, boolean streamable) {
this.name = name;
this.description = description;
this.inputSchema = inputSchema;
this.streamable = streamable;
}
public String getName() {
return name;
}
public String getDescription() {
return description;
}
public McpSchema.JsonSchema getInputSchema() {
return inputSchema;
}
public static Builder builder() {
return new Builder();
}
public static final class Builder {
private String name;
private String description;
private McpSchema.JsonSchema inputSchema;
private boolean streamable;
private Builder() {
}
public Builder name(String name) {
this.name = name;
return this;
}
public Builder description(String description) {
this.description = description;
return this;
}
public Builder inputSchema(McpSchema.JsonSchema inputSchema) {
this.inputSchema = inputSchema;
return this;
}
public Builder streamable(boolean streamable) {
this.streamable = streamable;
return this;
}
public ToolDefinition build() {
return new ToolDefinition(this.name, this.description, this.inputSchema, this.streamable);
}
}
}

View File

@ -0,0 +1,51 @@
package com.taobao.arthas.mcp.server.tool.definition;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.util.JsonSchemaGenerator;
import com.taobao.arthas.mcp.server.util.Assert;
import java.lang.reflect.Method;
public class ToolDefinitions {
public static ToolDefinition.Builder builder(Method method) {
Assert.notNull(method, "method cannot be null");
return ToolDefinition.builder()
.name(getToolName(method))
.description(getToolDescription(method))
.inputSchema(JsonSchemaGenerator.generateForMethodInput(method))
.streamable(isStreamable(method));
}
public static ToolDefinition from(Method method) {
return builder(method).build();
}
public static String getToolName(Method method) {
Assert.notNull(method, "method cannot be null");
Tool tool = method.getAnnotation(Tool.class);
if (tool == null) {
return method.getName();
}
return tool.name() != null ? tool.name() : method.getName();
}
public static String getToolDescription(Method method) {
Assert.notNull(method, "method cannot be null");
Tool tool = method.getAnnotation(Tool.class);
if (tool == null) {
return method.getName();
}
return tool.description() != null ? tool.description() : method.getName();
}
public static boolean isStreamable(Method method) {
Assert.notNull(method, "method cannot be null");
Tool tool = method.getAnnotation(Tool.class);
if (tool == null) {
return false;
}
return tool.streamable();
}
}

View File

@ -0,0 +1,76 @@
package com.taobao.arthas.mcp.server.tool.execution;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.util.JsonParser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.imageio.ImageIO;
import java.awt.image.RenderedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.lang.reflect.Type;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
/**
* A default implementation of {@link ToolCallResultConverter}.
*/
public final class DefaultToolCallResultConverter implements ToolCallResultConverter {
private static final Logger logger = LoggerFactory.getLogger(DefaultToolCallResultConverter.class);
private static final ObjectMapper OBJECT_MAPPER = JsonParser.getObjectMapper();
@Override
public String convert(Object result, Type returnType) {
if (returnType == Void.TYPE) {
logger.debug("The tool has no return type. Converting to conventional response.");
return JsonParser.toJson("Done");
}
if (result instanceof RenderedImage) {
final ByteArrayOutputStream buf = new ByteArrayOutputStream(1024 * 4);
try {
ImageIO.write((RenderedImage) result, "PNG", buf);
}
catch (IOException e) {
return "Failed to convert tool result to a base64 image: " + e.getMessage();
}
final String imgB64 = Base64.getEncoder().encodeToString(buf.toByteArray());
Map<String, String> imageData = new HashMap<>();
imageData.put("mimeType", "image/png");
imageData.put("data", imgB64);
return JsonParser.toJson(imageData);
}
else if (result instanceof String) {
String stringResult = (String) result;
if (isValidJson(stringResult)) {
logger.debug("Result is already valid JSON, returning as is.");
return stringResult;
} else {
logger.debug("Converting string result to JSON.");
return JsonParser.toJson(result);
}
}
else {
logger.debug("Converting tool result to JSON.");
return JsonParser.toJson(result);
}
}
private boolean isValidJson(String jsonString) {
if (jsonString == null || jsonString.trim().isEmpty()) {
return false;
}
try {
OBJECT_MAPPER.readTree(jsonString);
return true;
} catch (JsonProcessingException e) {
return false;
}
}
}

View File

@ -0,0 +1,52 @@
package com.taobao.arthas.mcp.server.tool.execution;
import com.taobao.arthas.mcp.server.util.Assert;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Default implementation of {@link ToolExecutionExceptionProcessor}.
*/
public class DefaultToolExecutionExceptionProcessor implements ToolExecutionExceptionProcessor {
private final static Logger logger = LoggerFactory.getLogger(DefaultToolExecutionExceptionProcessor.class);
private static final boolean DEFAULT_ALWAYS_THROW = false;
private final boolean alwaysThrow;
public DefaultToolExecutionExceptionProcessor(boolean alwaysThrow) {
this.alwaysThrow = alwaysThrow;
}
@Override
public String process(ToolExecutionException exception) {
Assert.notNull(exception, "exception cannot be null");
if (this.alwaysThrow) {
throw exception;
}
logger.debug("Exception thrown by tool: {}. Message: {}", exception.getToolDefinition().getName(),
exception.getMessage());
return exception.getMessage();
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private boolean alwaysThrow = DEFAULT_ALWAYS_THROW;
public Builder alwaysThrow(boolean alwaysThrow) {
this.alwaysThrow = alwaysThrow;
return this;
}
public DefaultToolExecutionExceptionProcessor build() {
return new DefaultToolExecutionExceptionProcessor(this.alwaysThrow);
}
}
}

View File

@ -0,0 +1,18 @@
package com.taobao.arthas.mcp.server.tool.execution;
import java.lang.reflect.Type;
/**
* A functional interface to convert tool call results to a String that can be sent back
* to the AI model.
*/
@FunctionalInterface
public interface ToolCallResultConverter {
/**
* Given an Object returned by a tool, convert it to a String compatible with the
* given class type.
*/
String convert(Object result, Type returnType);
}

View File

@ -0,0 +1,21 @@
package com.taobao.arthas.mcp.server.tool.execution;
import com.taobao.arthas.mcp.server.tool.definition.ToolDefinition;
/**
* An exception thrown when a tool execution fails.
*/
public class ToolExecutionException extends RuntimeException {
private final ToolDefinition toolDefinition;
public ToolExecutionException(ToolDefinition toolDefinition, Throwable cause) {
super(cause.getMessage(), cause);
this.toolDefinition = toolDefinition;
}
public ToolDefinition getToolDefinition() {
return this.toolDefinition;
}
}

View File

@ -0,0 +1,17 @@
package com.taobao.arthas.mcp.server.tool.execution;
/**
* A functional interface to process a {@link ToolExecutionException} by either converting
* the error message to a String that can be sent back to the AI model or throwing an
* exception to be handled by the caller.
*/
@FunctionalInterface
public interface ToolExecutionExceptionProcessor {
/**
* Convert an exception thrown by a tool to a String that can be sent back to the AI
* model or throw an exception to be handled by the caller.
*/
String process(ToolExecutionException exception);
}

View File

@ -0,0 +1,54 @@
package com.taobao.arthas.mcp.server.tool.function;
import com.taobao.arthas.mcp.server.ArthasMcpBootstrap;
import com.taobao.arthas.mcp.server.CommandExecutor;
import com.taobao.arthas.mcp.server.util.JsonParser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.HashMap;
import java.util.Map;
public class ArthasCommandExecutor {
private static final Logger logger = LoggerFactory.getLogger(ArthasCommandExecutor.class);
private static final long DEFAULT_TIMEOUT = 15000; // 默认超时时间15秒
/**
* Execute Arthas commands
* @param command command string
* @return Execution result as JSON string
*/
public static String executeCommand(String command) {
return executeCommand(command, DEFAULT_TIMEOUT);
}
/**
* Execute Arthas commands
* @param command command string
* @param timeout timeout (ms)
* @return Execution result as JSON string
*/
public static String executeCommand(String command, long timeout) {
try {
ArthasMcpBootstrap bootstrap = ArthasMcpBootstrap.getInstance();
if (bootstrap == null) {
throw new IllegalStateException("ArthasMcpBootstrap not initialized");
}
CommandExecutor executor = bootstrap.getCommandExecutor();
if (executor == null) {
throw new IllegalStateException("CommandExecutor not initialized");
}
Map<String, Object> result = executor.execute(command, timeout);
return JsonParser.toJson(result);
} catch (Exception e) {
logger.error("Failed to execute command: {}", command, e);
Map<String, Object> errorResponse = new HashMap<>();
errorResponse.put("error", true);
errorResponse.put("message", "Error executing command '" + command + "': " + e.getMessage());
return JsonParser.toJson(errorResponse);
}
}
}

View File

@ -0,0 +1,48 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class GetStaticTool {
@Tool(
name = "getstatic",
description = "GetStatic 诊断工具: 查看类的静态字段值,可指定 ClassLoader支持在返回结果上执行 OGNL 表达式。对应 Arthas 的 getstatic 命令。"
)
public String getstatic(
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass,
@ToolParam(description = "类名表达式匹配如java.lang.String或demo.MathGame")
String className,
@ToolParam(description = "静态字段名")
String fieldName,
@ToolParam(description = "OGNL 表达式", required = false)
String ognlExpression
) {
StringBuilder cmd = new StringBuilder("getstatic");
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
}
if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
cmd.append(" ").append(className.trim());
cmd.append(" ").append(fieldName.trim());
if (ognlExpression != null && !ognlExpression.trim().isEmpty()) {
cmd.append(" ").append(ognlExpression.trim());
}
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,16 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class JvmTool {
@Tool(
name = "jvm",
description = "Jvm 诊断工具: 查看当前 JVM 运行时信息。对应 Arthas 的 jvm 命令。"
)
public String jvm() {
String commandStr = "jvm";
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,16 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class MemoryTool {
@Tool(
name = "memory",
description = "Memory 诊断工具: 查看 JVM 内存使用情况,对应 Arthas 的 memory 命令。"
)
public String memory() {
String commandStr = "memory";
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,42 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class OgnlTool {
@Tool(
name = "ognl",
description = "OGNL 诊断工具: 执行 OGNL 表达式,对应 Arthas 的 ognl 命令。"
)
public String ognl(
@ToolParam(description = "OGNL 表达式")
String expression,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass,
@ToolParam(description = "结果对象展开层次 (-x),默认 1", required = false)
Integer expandLevel
) {
StringBuilder cmd = new StringBuilder("ognl");
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
}
if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
if (expandLevel != null && expandLevel > 0) {
cmd.append(" -x ").append(expandLevel);
}
cmd.append(" ").append(expression.trim());
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,24 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class PerfCounterTool {
@Tool(
name = "perfcounter",
description = "PerfCounter 诊断工具: 查看 JVM Perf Counter 信息,对应 Arthas 的 perfcounter 命令。"
)
public String perfcounter(
@ToolParam(description = "是否打印更多详情 (-d)", required = false)
Boolean detailed
) {
StringBuilder cmd = new StringBuilder("perfcounter");
if (Boolean.TRUE.equals(detailed)) {
cmd.append(" -d");
}
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,24 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class SysEnvTool {
@Tool(
name = "sysenv",
description = "SysEnv 诊断工具: 查看系统环境变量,对应 Arthas 的 sysenv 命令。"
)
public String sysenv(
@ToolParam(description = "环境变量名。若为空或空字符串,则查看所有变量;否则查看单个变量值。", required = false)
String envName
) {
StringBuilder cmd = new StringBuilder("sysenv");
if (envName != null && !envName.trim().isEmpty()) {
cmd.append(" ").append(envName.trim());
}
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,30 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class SysPropTool {
@Tool(
name = "sysprop",
description = "SysProp 诊断工具: 查看或修改系统属性,对应 Arthas 的 sysprop 命令。"
)
public String sysprop(
@ToolParam(description = "属性名", required = false)
String propertyName,
@ToolParam(description = "属性值;若指定则修改,否则查看", required = false)
String propertyValue
) {
StringBuilder cmd = new StringBuilder("sysprop");
if (propertyName != null && !propertyName.trim().isEmpty()) {
cmd.append(" ").append(propertyName.trim());
if (propertyValue != null && !propertyValue.trim().isEmpty()) {
cmd.append(" ").append(propertyValue.trim());
}
}
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,30 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class VMOptionTool {
@Tool(
name = "vmoption",
description = "VMOption 诊断工具: 查看或更新 JVM VM options对应 Arthas 的 vmoption 命令。"
)
public String vmoption(
@ToolParam(description = "Name of the VM option.", required = false)
String key,
@ToolParam(description = "更新值,仅在更新时使用", required = false)
String value
) {
StringBuilder cmd = new StringBuilder("vmoption");
if (key != null && !key.trim().isEmpty()) {
cmd.append(" ").append(key.trim());
if (value != null && !value.trim().isEmpty()) {
cmd.append(" ").append(value.trim());
}
}
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,83 @@
package com.taobao.arthas.mcp.server.tool.function.jvm300;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class VMToolTool {
public static final String ACTION_GET_INSTANCES = "getInstances";
public static final String ACTION_FORCE_GC = "forceGc";
public static final String ACTION_INTERRUPT_THREAD = "interruptThread";
@Tool(
name = "vmtool",
description = "虚拟机工具诊断工具: 查询实例、强制 GC、线程中断等对应 Arthas 的 vmtool 命令。"
)
public String vmtool(
@ToolParam(description = "操作类型: getInstances/forceGc/interruptThread 等")
String action,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass,
@ToolParam(description = "类名全限定getInstances 时使用)", required = false)
String className,
@ToolParam(description = "返回实例限制数量 (-l)getInstances 时使用,默认 10<=0 表示不限制", required = false)
Integer limit,
@ToolParam(description = "结果对象展开层次 (-x),默认 1", required = false)
Integer expandLevel,
@ToolParam(description = "OGNL 表达式,对 getInstances 返回的 instances 执行 (--express)", required = false)
String express,
@ToolParam(description = "线程 ID (-t)interruptThread 时使用", required = false)
Long threadId
) {
StringBuilder cmd = new StringBuilder("vmtool");
if (action == null || action.trim().isEmpty()) {
throw new IllegalArgumentException("vmtool: action 参数不能为空");
}
cmd.append(" --action ").append(action.trim());
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
}
if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
if (ACTION_GET_INSTANCES.equals(action.trim())) {
if (className != null && !className.trim().isEmpty()) {
cmd.append(" --className ").append(className.trim());
}
if (limit != null) {
cmd.append(" --limit ").append(limit);
}
if (expandLevel != null && expandLevel > 0) {
cmd.append(" -x ").append(expandLevel);
}
if (express != null && !express.trim().isEmpty()) {
cmd.append(" --express ").append(express.trim());
}
}
// interruptThread
if (ACTION_INTERRUPT_THREAD.equals(action.trim())) {
if (threadId != null && threadId > 0) {
cmd.append(" -t ").append(threadId);
} else {
throw new IllegalArgumentException("vmtool interruptThread 需要指定线程 ID (threadId)");
}
}
String commandStr = cmd.toString();
return ArthasCommandExecutor.executeCommand(commandStr);
}
}

View File

@ -0,0 +1,73 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class ClassLoaderTool {
public static final String MODE_STATS = "stats";
public static final String MODE_INSTANCES = "instances";
public static final String MODE_TREE = "tree";
public static final String MODE_ALL_CLASSES = "all-classes";
public static final String MODE_URL_STATS = "url-stats";
@Tool(
name = "classloader",
description = "ClassLoader 诊断工具可以查看类加载器统计信息、继承树、URLs以及进行资源查找和类加载操作"
)
public String classloader(
@ToolParam(description = "显示模式stats(统计信息,默认), instances(实例详情), tree(继承树), all-classes(所有类,慎用), url-stats(URL统计)", required = false)
String mode,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass,
@ToolParam(description = "要查找的资源名称如META-INF/MANIFEST.MF", required = false)
String resource,
@ToolParam(description = "要加载的类名,支持全限定名", required = false)
String loadClass) {
StringBuilder cmd = new StringBuilder("classloader");
if (mode != null) {
switch (mode.toLowerCase()) {
case MODE_INSTANCES:
cmd.append(" -l");
break;
case MODE_TREE:
cmd.append(" -t");
break;
case MODE_ALL_CLASSES:
cmd.append(" -a");
break;
case MODE_URL_STATS:
cmd.append(" --url-stat");
break;
case MODE_STATS:
default:
break;
}
}
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
} else if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
if (resource != null && !resource.trim().isEmpty()) {
cmd.append(" -r ").append(resource.trim());
}
if (loadClass != null && !loadClass.trim().isEmpty()) {
cmd.append(" --load ").append(loadClass.trim());
}
return ArthasCommandExecutor.executeCommand(cmd.toString());
}
}

View File

@ -0,0 +1,63 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class JadTool {
@Tool(
name = "jad",
description = "反编译指定已加载类的源码将JVM中实际运行的class的bytecode反编译成java代码"
)
public String jad(
@ToolParam(description = "类名表达式匹配如java.lang.String或demo.MathGame")
String classPattern,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass,
@ToolParam(description = "反编译时只显示源代码默认false", required = false)
Boolean sourceOnly,
@ToolParam(description = "反编译时不显示行号默认false", required = false)
Boolean noLineNumber,
@ToolParam(description = "开启正则表达式匹配默认为通配符匹配默认false", required = false)
Boolean useRegex,
@ToolParam(description = "指定dump class文件目录默认会dump到logback.xml中配置的log目录下", required = false)
String dumpDirectory) {
StringBuilder cmd = new StringBuilder("jad");
cmd.append(" ").append(classPattern);
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
} else if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
if (Boolean.TRUE.equals(sourceOnly)) {
cmd.append(" --source-only");
}
if (Boolean.TRUE.equals(noLineNumber)) {
cmd.append(" --lineNumber false");
}
if (Boolean.TRUE.equals(useRegex)) {
cmd.append(" -E");
}
if (dumpDirectory != null && !dumpDirectory.trim().isEmpty()) {
cmd.append(" -d ").append(dumpDirectory.trim());
}
return ArthasCommandExecutor.executeCommand(cmd.toString());
}
}

View File

@ -0,0 +1,42 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class MemoryCompilerTool {
@Tool(
name = "mc",
description = "Memory Compiler/内存编译器,编译.java文件生成.class"
)
public String mc(
@ToolParam(description = "要编译的.java文件路径支持多个文件用空格分隔")
String javaFilePaths,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass,
@ToolParam(description = "指定输出目录", required = false)
String outputDir) {
StringBuilder cmd = new StringBuilder("mc");
cmd.append(" ").append(javaFilePaths);
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
} else if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
if (outputDir != null && !outputDir.trim().isEmpty()) {
cmd.append(" -d ").append(outputDir.trim());
}
return ArthasCommandExecutor.executeCommand(cmd.toString());
}
}

View File

@ -0,0 +1,35 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class RedefineTool {
@Tool(
name = "redefine",
description = "重新加载类的字节码允许在JVM运行时重新加载已存在的类的字节码实现热更新"
)
public String redefine(
@ToolParam(description = "要重新定义的.class文件路径支持多个文件用空格分隔")
String classFilePaths,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "指定执行表达式的ClassLoader的class name如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass) {
StringBuilder cmd = new StringBuilder("redefine");
cmd.append(" ").append(classFilePaths);
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
} else if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
return ArthasCommandExecutor.executeCommand(cmd.toString());
}
}

View File

@ -0,0 +1,35 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class RetransformTool {
@Tool(
name = "retransform",
description = "热加载类的字节码,允许对已加载的类进行字节码修改并使其生效"
)
public String retransform(
@ToolParam(description = "要操作的.class文件路径支持多个文件用空格分隔")
String classFilePaths,
@ToolParam(description = "ClassLoader的hashcode16进制用于指定特定的ClassLoader", required = false)
String classLoaderHash,
@ToolParam(description = "ClassLoader的完整类名如sun.misc.Launcher$AppClassLoader可替代hashcode", required = false)
String classLoaderClass) {
StringBuilder cmd = new StringBuilder("retransform");
cmd.append(" ").append(classFilePaths);
if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) {
cmd.append(" -c ").append(classLoaderHash.trim());
} else if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) {
cmd.append(" --classLoaderClass ").append(classLoaderClass.trim());
}
return ArthasCommandExecutor.executeCommand(cmd.toString());
}
}

View File

@ -0,0 +1,19 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class SearchClassTool {
@Tool(
name = "getClassInfo",
description = "获取指定类的详细信息,包括类加载器、方法、字段等信息"
)
public String sc(
@ToolParam(description = "要查询的类名,支持全限定名") String className) {
return ArthasCommandExecutor.executeCommand("sc -d " + className);
}
}

View File

@ -0,0 +1,18 @@
package com.taobao.arthas.mcp.server.tool.function.klass100;
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor;
public class SearchMethodTool {
@Tool(
name = "getClassMethods",
description = "获取指定类的所有方法信息"
)
public String sm(
@ToolParam(description = "要查询的类名,支持全限定名") String className) {
return ArthasCommandExecutor.executeCommand("sm " + className);
}
}

View File

@ -0,0 +1,174 @@
package com.taobao.arthas.mcp.server.tool.util;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
import com.taobao.arthas.mcp.server.util.Assert;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* Simple JsonSchema generator
* JsonSchema definitions for generating method parameters
*/
public final class JsonSchemaGenerator {
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private static final boolean PROPERTY_REQUIRED_BY_DEFAULT = true;
private JsonSchemaGenerator() {
}
/**
* Generate JsonSchema for method parameters
* @param method target method
* @return JsonSchema object
*/
public static McpSchema.JsonSchema generateForMethodInput(Method method) {
Assert.notNull(method, "method cannot be null");
ObjectNode schema = OBJECT_MAPPER.createObjectNode();
schema.put("type", "object");
ObjectNode properties = schema.putObject("properties");
List<String> required = new ArrayList<>();
Parameter[] parameters = method.getParameters();
for (Parameter parameter : parameters) {
String paramName = getParameterName(parameter);
Class<?> paramType = parameter.getType();
boolean isRequired = isParameterRequired(parameter);
if (isRequired) {
required.add(paramName);
}
ObjectNode paramProperties = generateParameterProperties(paramType);
String description = getParameterDescription(parameter);
if (description != null) {
paramProperties.put("description", description);
}
properties.set(paramName, paramProperties);
}
if (!required.isEmpty()) {
ArrayNode requiredArray = schema.putArray("required");
for (String req : required) {
requiredArray.add(req);
}
}
schema.put("additionalProperties", false);
return new McpSchema.JsonSchema("object", convertToMap(properties), required, false);
}
private static String getParameterName(Parameter parameter) {
JsonProperty jsonProperty = parameter.getAnnotation(JsonProperty.class);
if (jsonProperty != null && !jsonProperty.value().isEmpty()) {
return jsonProperty.value();
}
return parameter.getName();
}
private static ObjectNode generateParameterProperties(Class<?> paramType) {
ObjectNode properties = OBJECT_MAPPER.createObjectNode();
if (paramType == String.class) {
properties.put("type", "string");
} else if (paramType == int.class || paramType == Integer.class
|| paramType == long.class || paramType == Long.class) {
properties.put("type", "integer");
} else if (paramType == double.class || paramType == Double.class
|| paramType == float.class || paramType == Float.class) {
properties.put("type", "number");
} else if (paramType == boolean.class || paramType == Boolean.class) {
properties.put("type", "boolean");
} else if (paramType.isArray()) {
properties.put("type", "array");
ObjectNode items = properties.putObject("items");
Class<?> componentType = paramType.getComponentType();
if (componentType == String.class) {
items.put("type", "string");
} else if (componentType == int.class || componentType == Integer.class
|| componentType == long.class || componentType == Long.class) {
items.put("type", "integer");
} else if (componentType == double.class || componentType == Double.class
|| componentType == float.class || componentType == Float.class) {
items.put("type", "number");
} else if (componentType == boolean.class || componentType == Boolean.class) {
items.put("type", "boolean");
} else {
items.put("type", "object");
}
} else {
properties.put("type", "object");
}
return properties;
}
private static boolean isParameterRequired(Parameter parameter) {
ToolParam toolParam = parameter.getAnnotation(ToolParam.class);
if (toolParam != null) {
return toolParam.required();
}
JsonProperty jsonProperty = parameter.getAnnotation(JsonProperty.class);
if (jsonProperty != null) {
return jsonProperty.required();
}
return PROPERTY_REQUIRED_BY_DEFAULT;
}
private static String getParameterDescription(Parameter parameter) {
ToolParam toolParam = parameter.getAnnotation(ToolParam.class);
if (toolParam != null && toolParam.description() != null && !toolParam.description().isEmpty()) {
return toolParam.description();
}
JsonPropertyDescription jsonPropertyDescription = parameter.getAnnotation(JsonPropertyDescription.class);
if (jsonPropertyDescription != null && !jsonPropertyDescription.value().isEmpty()) {
return jsonPropertyDescription.value();
}
return null;
}
private static Map<String, Object> convertToMap(ObjectNode node) {
Map<String, Object> result = new HashMap<>();
node.fields().forEachRemaining(entry -> {
JsonNode value = entry.getValue();
if (value.isObject()) {
result.put(entry.getKey(), convertToMap((ObjectNode) value));
} else if (value.isArray()) {
List<Object> array = new ArrayList<>();
value.elements().forEachRemaining(element -> {
if (element.isObject()) {
array.add(convertToMap((ObjectNode) element));
} else {
array.add(element.asText());
}
});
result.put(entry.getKey(), array);
} else {
result.put(entry.getKey(), value.asText());
}
});
return result;
}
}

View File

@ -0,0 +1,85 @@
package com.taobao.arthas.mcp.server.tool.util;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.taobao.arthas.mcp.server.protocol.config.McpServerProperties;
import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange;
import com.taobao.arthas.mcp.server.protocol.server.McpServerFeatures;
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
import com.taobao.arthas.mcp.server.tool.ToolCallback;
import com.taobao.arthas.mcp.server.tool.ToolContext;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
public final class McpToolUtils {
public static final String TOOL_CONTEXT_MCP_EXCHANGE_KEY = "exchange";
private McpToolUtils() {
}
public static List<McpServerFeatures.ToolSpecification> toToolSpecifications(
List<ToolCallback> tools,
McpServerProperties serverProperties) {
// De-duplicate tools by their name, keeping the first occurrence of each tool name
return tools.stream()
.collect(Collectors.toMap(
tool -> tool.getToolDefinition().getName(), // Key: tool name
tool -> tool, // Value: the tool itself
(existing, replacement) -> existing // On duplicate key, keep the existing tool
))
.values()
.stream()
.map(McpToolUtils::toToolSpecification)
.collect(Collectors.toList());
}
public static McpServerFeatures.ToolSpecification toToolSpecification(ToolCallback toolCallback) {
McpSchema.Tool tool = new McpSchema.Tool(toolCallback.getToolDefinition().getName(),
toolCallback.getToolDefinition().getDescription(), toolCallback.getToolDefinition().getInputSchema());
return new McpServerFeatures.ToolSpecification(tool, new BiFunction<McpNettyServerExchange, Map<String, Object>, CompletableFuture<McpSchema.CallToolResult>>() {
@Override
public CompletableFuture<McpSchema.CallToolResult> apply(McpNettyServerExchange exchange, Map<String, Object> request) {
try {
Map<String, Object> contextMap = new HashMap<>();
contextMap.put(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchange);
ToolContext toolContext = new ToolContext(contextMap);
String callResult = toolCallback.call(convertRequestToString(request), toolContext);
List<McpSchema.Content> contents = new ArrayList<>();
contents.add(new McpSchema.TextContent(callResult));
return CompletableFuture.completedFuture(new McpSchema.CallToolResult(contents, false));
}
catch (Exception e) {
List<McpSchema.Content> contents = new ArrayList<>();
contents.add(new McpSchema.TextContent(e.getMessage()));
return CompletableFuture.completedFuture(new McpSchema.CallToolResult(contents, true));
}
}
private String convertRequestToString(Map<String, Object> request) {
if (request == null) {
return "";
}
try {
return new ObjectMapper().writeValueAsString(request);
} catch (Exception e) {
return request.toString();
}
}
});
}
}

View File

@ -0,0 +1,44 @@
package com.taobao.arthas.mcp.server.util;
import java.util.Collection;
import java.util.Map;
/**
* Assertion utility class for parameter validation.
*/
public final class Assert {
private Assert() {
}
public static void notNull(Object object, String message) {
if (object == null) {
throw new IllegalArgumentException(message);
}
}
public static void hasText(String text, String message) {
if (text == null || text.trim().isEmpty()) {
throw new IllegalArgumentException(message);
}
}
public static void notEmpty(Collection<?> collection, String message) {
if (collection == null || collection.isEmpty()) {
throw new IllegalArgumentException(message);
}
}
public static void notEmpty(Map<?, ?> map, String message) {
if (map == null || map.isEmpty()) {
throw new IllegalArgumentException(message);
}
}
public static void isTrue(boolean condition, String message) {
if (!condition) {
throw new IllegalArgumentException(message);
}
}
}

View File

@ -0,0 +1,141 @@
package com.taobao.arthas.mcp.server.util;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.databind.json.JsonMapper;
import java.lang.reflect.Type;
import java.math.BigDecimal;
/**
* Utilities to perform parsing operations between JSON and Java.
*/
public final class JsonParser {
private static final ObjectMapper OBJECT_MAPPER = JsonMapper.builder()
.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
.disable(SerializationFeature.FAIL_ON_EMPTY_BEANS)
.build();
private JsonParser() {
}
public static ObjectMapper getObjectMapper() {
return OBJECT_MAPPER;
}
public static <T> T fromJson(String json, Class<T> type) {
Assert.notNull(json, "json cannot be null");
Assert.notNull(type, "type cannot be null");
try {
return OBJECT_MAPPER.readValue(json, type);
}
catch (JsonProcessingException ex) {
throw new IllegalStateException("Conversion from JSON to " + type.getName() + " failed", ex);
}
}
public static <T> T fromJson(String json, Type type) {
Assert.notNull(json, "json cannot be null");
Assert.notNull(type, "type cannot be null");
try {
return OBJECT_MAPPER.readValue(json, OBJECT_MAPPER.constructType(type));
}
catch (JsonProcessingException ex) {
throw new IllegalStateException("Conversion from JSON to " + type.getTypeName() + " failed", ex);
}
}
public static <T> T fromJson(String json, TypeReference<T> type) {
Assert.notNull(json, "json cannot be null");
Assert.notNull(type, "type cannot be null");
try {
return OBJECT_MAPPER.readValue(json, type);
}
catch (JsonProcessingException ex) {
throw new IllegalStateException("Conversion from JSON to " + type.getType().getTypeName() + " failed",
ex);
}
}
/**
* Converts a Java object to a JSON string.
*/
public static String toJson(Object object) {
try {
return OBJECT_MAPPER.writeValueAsString(object);
}
catch (JsonProcessingException ex) {
throw new IllegalStateException("Conversion from Object to JSON failed", ex);
}
}
public static Object toTypedObject(Object value, Class<?> type) {
if (value == null) {
throw new IllegalArgumentException("value cannot be null");
}
if (type == null) {
throw new IllegalArgumentException("type cannot be null");
}
Class<?> javaType = resolvePrimitiveIfNecessary(type);
if (javaType == String.class) {
return value.toString();
}
else if (javaType == Byte.class) {
return Byte.parseByte(value.toString());
}
else if (javaType == Integer.class) {
BigDecimal bigDecimal = new BigDecimal(value.toString());
return bigDecimal.intValueExact();
}
else if (javaType == Short.class) {
return Short.parseShort(value.toString());
}
else if (javaType == Long.class) {
BigDecimal bigDecimal = new BigDecimal(value.toString());
return bigDecimal.longValueExact();
}
else if (javaType == Double.class) {
return Double.parseDouble(value.toString());
}
else if (javaType == Float.class) {
return Float.parseFloat(value.toString());
}
else if (javaType == Boolean.class) {
return Boolean.parseBoolean(value.toString());
}
else if (javaType == Character.class) {
String s = value.toString();
if (s.length() == 1) {
return s.charAt(0);
}
throw new IllegalArgumentException("Cannot convert to char: " + value);
}
else if (javaType.isEnum()) {
@SuppressWarnings("unchecked")
Class<Enum> enumType = (Class<Enum>) javaType;
return Enum.valueOf(enumType, value.toString());
}
String json = JsonParser.toJson(value);
return JsonParser.fromJson(json, javaType);
}
public static Class<?> resolvePrimitiveIfNecessary(Class<?> type) {
if (type.isPrimitive()) {
return Utils.getWrapperClassForPrimitive(type);
}
return type;
}
}

View File

@ -0,0 +1,71 @@
package com.taobao.arthas.mcp.server.util;
import java.util.Collection;
import java.util.Map;
public final class Utils {
public static boolean hasText(String str) {
return str != null && !str.trim().isEmpty();
}
public static boolean isEmpty(Collection<?> collection) {
return (collection == null || collection.isEmpty());
}
public static boolean isEmpty(Map<?, ?> map) {
return (map == null || map.isEmpty());
}
public static boolean isAssignable(Class<?> targetType, Class<?> sourceType) {
if (targetType == null || sourceType == null) {
return false;
}
if (targetType.equals(sourceType)) {
return true;
}
if (targetType.isAssignableFrom(sourceType)) {
return true;
}
if (targetType.isPrimitive()) {
Class<?> resolvedPrimitive = getPrimitiveClassForWrapper(sourceType);
return resolvedPrimitive != null && targetType.equals(resolvedPrimitive);
}
else if (sourceType.isPrimitive()) {
Class<?> resolvedWrapper = getWrapperClassForPrimitive(sourceType);
return resolvedWrapper != null && targetType.equals(resolvedWrapper);
}
return false;
}
public static Class<?> getPrimitiveClassForWrapper(Class<?> wrapperClass) {
if (Boolean.class.equals(wrapperClass)) return boolean.class;
if (Byte.class.equals(wrapperClass)) return byte.class;
if (Character.class.equals(wrapperClass)) return char.class;
if (Double.class.equals(wrapperClass)) return double.class;
if (Float.class.equals(wrapperClass)) return float.class;
if (Integer.class.equals(wrapperClass)) return int.class;
if (Long.class.equals(wrapperClass)) return long.class;
if (Short.class.equals(wrapperClass)) return short.class;
if (Void.class.equals(wrapperClass)) return void.class;
return null;
}
public static Class<?> getWrapperClassForPrimitive(Class<?> primitiveClass) {
if (boolean.class.equals(primitiveClass)) return Boolean.class;
if (byte.class.equals(primitiveClass)) return Byte.class;
if (char.class.equals(primitiveClass)) return Character.class;
if (double.class.equals(primitiveClass)) return Double.class;
if (float.class.equals(primitiveClass)) return Float.class;
if (int.class.equals(primitiveClass)) return Integer.class;
if (long.class.equals(primitiveClass)) return Long.class;
if (short.class.equals(primitiveClass)) return Short.class;
if (void.class.equals(primitiveClass)) return Void.class;
return null;
}
}

View File

@ -81,6 +81,7 @@
<module>labs/cluster-management/native-agent-proxy</module> <module>labs/cluster-management/native-agent-proxy</module>
<module>labs/cluster-management/native-agent-common</module> <module>labs/cluster-management/native-agent-common</module>
<module>labs/arthas-grpc-server</module> <module>labs/arthas-grpc-server</module>
<module>labs/arthas-mcp-server</module>
</modules> </modules>
<properties> <properties>