From 0b616ece0e2bcb075449e992b27f003eb19f10fe Mon Sep 17 00:00:00 2001 From: Zhang Yu <143883582+Yeaury@users.noreply.github.com> Date: Fri, 29 Aug 2025 23:10:28 +0800 Subject: [PATCH] feat: Support MCP 2025-03-26 protocol with streaming tools (#3054) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 实现MCP协议2025-03-26版本支持 - 添加流式工具(Streamable Tools)架构 - 支持双传输模式：流式传输(SSE)和无状态传输 - 实现30+个Arthas工具的MCP接口封装 - 添加实时进度通知和结果流式输出 - 支持watch、monitor、trace、dashboard、profiler等流式工具 --- .../core/command/CommandExecutorImpl.java | 463 ++++++++---- .../arthas/core/server/ArthasBootstrap.java | 7 +- .../term/impl/http/HttpRequestHandler.java | 9 +- labs/arthas-mcp-server/pom.xml | 7 +- .../arthas/mcp/server/ArthasMcpBootstrap.java | 17 +- .../arthas/mcp/server/ArthasMcpServer.java | 173 +++-- .../arthas/mcp/server/CommandExecutor.java | 23 +- .../protocol/config/McpServerProperties.java | 47 +- .../DefaultMcpStatelessServerHandler.java | 85 +++ .../server/DefaultMcpTransportContext.java | 47 ++ .../server/McpInitRequestHandler.java | 24 + .../protocol/server/McpNettyServer.java | 134 ++-- .../server/McpNettyServerExchange.java | 110 ++- .../server/McpNotificationHandler.java | 26 + .../protocol/server/McpRequestHandler.java | 26 + .../mcp/server/protocol/server/McpServer.java | 319 +++++--- .../protocol/server/McpServerFeatures.java | 31 +- .../server/McpStatelessNettyServer.java | 456 ++++++++++++ .../McpStatelessNotificationHandler.java | 23 + .../server/McpStatelessRequestHandler.java | 24 + .../server/McpStatelessServerFeatures.java | 218 ++++++ .../server/McpStatelessServerHandler.java | 32 + .../protocol/server/McpTransportContext.java | 19 + .../server/McpTransportContextExtractor.java | 25 + .../server/handler/McpHttpRequestHandler.java | 324 +++++++++ .../server/handler/McpRequestHandler.java | 379 ---------- .../McpStatelessHttpRequestHandler.java | 227 ++++++ .../McpStreamableHttpRequestHandler.java | 685 ++++++++++++++++++ .../server/store/InMemoryEventStore.java | 155 ++++ .../HttpNettyServerTransportProvider.java | 90 --- .../NettyStatelessServerTransport.java | 112 +++ ...ettyStreamableServerTransportProvider.java | 148 ++++ ...aultMcpStreamableServerSessionFactory.java | 64 ++ .../mcp/server/protocol/spec/EventStore.java | 51 ++ .../mcp/server/protocol/spec/HttpHeaders.java | 24 + .../mcp/server/protocol/spec/McpError.java | 4 + .../mcp/server/protocol/spec/McpSchema.java | 121 +++- .../protocol/spec/McpServerSession.java | 433 ----------- .../protocol/spec/McpServerTransport.java | 4 + .../spec/McpServerTransportProvider.java | 12 +- .../mcp/server/protocol/spec/McpSession.java | 4 + .../spec/McpStatelessServerTransport.java | 21 + .../spec/McpStreamableServerSession.java | 428 +++++++++++ .../spec/McpStreamableServerTransport.java | 21 + .../McpStreamableServerTransportProvider.java | 22 + .../server/protocol/spec/McpTransport.java | 4 + .../spec/MissingMcpTransportSession.java | 52 ++ .../server/session/ArthasCommandContext.java | 76 ++ .../session/ArthasCommandSessionManager.java | 137 ++++ .../mcp/server/tool/DefaultToolCallback.java | 32 + .../tool/definition/ToolDefinition.java | 4 + .../tool/function/ArthasCommandExecutor.java | 54 -- .../tool/function/StreamableToolUtils.java | 262 +++++++ .../tool/function/jvm300/DashboardTool.java | 81 +++ .../tool/function/jvm300/GetStaticTool.java | 13 +- .../tool/function/jvm300/HeapdumpTool.java | 70 ++ .../server/tool/function/jvm300/JvmTool.java | 11 +- .../tool/function/jvm300/MBeanTool.java | 123 ++++ .../tool/function/jvm300/MemoryTool.java | 11 +- .../server/tool/function/jvm300/OgnlTool.java | 13 +- .../tool/function/jvm300/PerfCounterTool.java | 12 +- .../tool/function/jvm300/SysEnvTool.java | 12 +- .../tool/function/jvm300/SysPropTool.java | 13 +- .../tool/function/jvm300/ThreadTool.java | 68 ++ .../tool/function/jvm300/VMOptionTool.java | 13 +- .../tool/function/jvm300/VMToolTool.java | 13 +- .../function/klass100/ClassLoaderTool.java | 13 +- .../tool/function/klass100/DumpClassTool.java | 68 ++ .../tool/function/klass100/JadTool.java | 13 +- .../function/klass100/MemoryCompilerTool.java | 21 +- .../tool/function/klass100/RedefineTool.java | 13 +- .../function/klass100/RetransformTool.java | 13 +- .../function/klass100/SearchClassTool.java | 12 +- .../function/klass100/SearchMethodTool.java | 12 +- .../tool/function/monitor200/MonitorTool.java | 122 ++++ .../tool/function/monitor200/StackTool.java | 106 +++ .../function/monitor200/TimeTunnelTool.java | 258 +++++++ .../tool/function/monitor200/TraceTool.java | 105 +++ .../tool/function/monitor200/WatchTool.java | 138 ++++ .../server/tool/util/JsonSchemaGenerator.java | 5 + .../mcp/server/tool/util/McpToolUtils.java | 173 +++-- .../arthas/mcp/server/util/JsonParser.java | 3 + .../mcp/server/util/KeepAliveScheduler.java | 209 ++++++ site/docs/doc/monitor.md | 2 +- site/docs/en/doc/monitor.md | 2 +- 85 files changed, 6496 insertions(+), 1540 deletions(-) create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpStatelessServerHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpTransportContext.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpInitRequestHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNotificationHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpRequestHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNettyServer.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNotificationHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessRequestHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerFeatures.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContext.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContextExtractor.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpHttpRequestHandler.java delete mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpRequestHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStatelessHttpRequestHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStreamableHttpRequestHandler.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/store/InMemoryEventStore.java delete mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/HttpNettyServerTransportProvider.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStatelessServerTransport.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStreamableServerTransportProvider.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/DefaultMcpStreamableServerSessionFactory.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/EventStore.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/HttpHeaders.java delete mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerSession.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStatelessServerTransport.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerSession.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransport.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransportProvider.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/MissingMcpTransportSession.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandContext.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java delete mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/ArthasCommandExecutor.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/StreamableToolUtils.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/DashboardTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/HeapdumpTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MBeanTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/ThreadTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/DumpClassTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/MonitorTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/StackTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TimeTunnelTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TraceTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/WatchTool.java create mode 100644 labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/KeepAliveScheduler.java diff --git a/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java b/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java index 0bd48d333..30044fc2c 100644 --- a/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java +++ b/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java @@ -1,9 +1,13 @@ 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.common.PidUtils; +import com.taobao.arthas.core.command.model.*; +import com.taobao.arthas.core.distribution.ResultConsumer; +import com.taobao.arthas.core.distribution.ResultDistributor; +import com.taobao.arthas.core.distribution.SharingResultDistributor; import com.taobao.arthas.core.distribution.impl.PackingResultDistributorImpl; +import com.taobao.arthas.core.distribution.impl.ResultConsumerImpl; +import com.taobao.arthas.core.distribution.impl.SharingResultDistributorImpl; import com.taobao.arthas.core.shell.cli.CliToken; import com.taobao.arthas.core.shell.cli.CliTokens; import com.taobao.arthas.core.shell.cli.Completion; @@ -16,24 +20,24 @@ 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.core.util.ArthasBanner; +import com.taobao.arthas.core.util.DateUtils; +import com.taobao.arthas.core.util.StringUtils; import com.taobao.arthas.mcp.server.CommandExecutor; +import io.termd.core.function.Function; 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命令 + * 命令执行器，用于执行Arthas命令，支持同步和异步执行 */ public class CommandExecutorImpl implements CommandExecutor { private static final Logger logger = LoggerFactory.getLogger(CommandExecutorImpl.class); + private static final String ONETIME_SESSION_KEY = "oneTimeSession"; private final SessionManager sessionManager; private final JobController jobController; @@ -45,42 +49,43 @@ public class CommandExecutorImpl implements CommandExecutor { this.jobController = sessionManager.getJobController(); } - /** - * 同步执行命令 - * @param commandLine 命令行 - * @param timeout 超时时间(毫秒) - * @return 命令执行结果 - */ - @Override - public Map 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 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 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"); + public Session getCurrentSession(String sessionId, boolean oneTimeIsAllowed) { + if (sessionId == null || sessionId.trim().isEmpty()) { + if (!oneTimeIsAllowed) { + throw new SessionNotFoundException("SessionId is required for this operation"); } - resultDistributor = new PackingResultDistributorImpl(session); - InMemoryTerm term = new InMemoryTerm(); - term.setSession(session); - job = jobController.createJob(commandManager, tokens, session, - new JobHandle(), term, resultDistributor); + Session session = sessionManager.createSession(); + if (session == null) { + throw new SessionNotFoundException("Failed to create temporary session"); + } + session.put(ONETIME_SESSION_KEY, new Object()); + logger.debug("Created one-time session {}", session.getSessionId()); + return session; + } else { + Session session = sessionManager.getSession(sessionId); + if (session == null) { + throw new SessionNotFoundException("Session not found: " + sessionId); + } + sessionManager.updateAccessTime(session); + logger.debug("Using existing session {}", sessionId); + return session; + } + } + + @Override + public Map executeSync(String commandLine, long timeout, String sessionId) { + Session session = null; + boolean oneTimeAccess = false; + + try { + session = getCurrentSession(sessionId, true); + if (session.get(ONETIME_SESSION_KEY) != null) { + oneTimeAccess = true; + } + + PackingResultDistributorImpl resultDistributor = new PackingResultDistributorImpl(session); + Job job = this.createJob(commandLine, session, resultDistributor); if (job == null) { logger.error("Failed to create job for command: {}", commandLine); @@ -88,7 +93,6 @@ public class CommandExecutorImpl implements CommandExecutor { } job.run(); - boolean finished = waitForJob(job, (int) timeout); if (!finished) { logger.warn("Command timeout after {} ms: {}", timeout, commandLine); @@ -96,11 +100,13 @@ public class CommandExecutorImpl implements CommandExecutor { return createTimeoutResult(commandLine, timeout); } + Map result = new TreeMap<>(); result.put("command", commandLine); result.put("success", true); + result.put("sessionId", session.getSessionId()); result.put("executionTime", System.currentTimeMillis()); - List results = resultDistributor.getResults(); + List results = resultDistributor.getResults(); if (results != null && !results.isEmpty()) { result.put("results", results); result.put("resultCount", results.size()); @@ -108,28 +114,198 @@ public class CommandExecutorImpl implements CommandExecutor { 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 (SessionNotFoundException e) { + logger.error("Session error for command: {}", commandLine, e); + return createErrorResult(commandLine, e.getMessage()); } catch (Exception e) { logger.error("Error executing command: {}", commandLine, e); return createErrorResult(commandLine, "Error executing command: " + e.getMessage()); } finally { - if (session != null) { + if (oneTimeAccess && session != null) { try { sessionManager.removeSession(session.getSessionId()); + logger.debug("Destroyed one-time session {}", session.getSessionId()); } catch (Exception e) { - logger.warn("Error removing session", e); + logger.warn("Error removing one-time session", e); } } } } + @Override + public Map executeAsync(String commandLine, String sessionId) { + Map result = new TreeMap<>(); + Session session = getCurrentSession(sessionId, false); + if (!session.tryLock()) { + logger.warn("Another command is executing in session: {}", session.getSessionId()); + return createErrorResult(commandLine, "Another command is executing"); + } + int lock = session.getLock(); + + try { + Job foregroundJob = session.getForegroundJob(); + if (foregroundJob != null) { + logger.warn("Another job is running in session: {}, jobId: {}", session.getSessionId(), foregroundJob.id()); + session.unLock(); + return createErrorResult(commandLine, "Another job is running, jobId: " + foregroundJob.id()); + } + + Job job = this.createJob(commandLine, session, session.getResultDistributor()); + + if (job == null) { + logger.error("Failed to create job for command: {}", commandLine); + session.unLock(); + return createErrorResult(commandLine, "Failed to create job"); + } + + session.setForegroundJob(job); + updateSessionInputStatus(session, InputStatus.ALLOW_INTERRUPT); + + job.run(); + + result.put("success", true); + result.put("command", commandLine); + result.put("sessionId", session.getSessionId()); + result.put("jobId", job.id()); + result.put("jobStatus", job.status().toString()); + + return result; + + } catch (SessionNotFoundException e) { + logger.error("Session error for async command: {}", commandLine, e); + return createErrorResult(commandLine, e.getMessage()); + } catch (Exception e) { + logger.error("Error executing async command: {}", commandLine, e); + return createErrorResult(commandLine, "Error executing async command: " + e.getMessage()); + }finally { + if (session.getLock() == lock) { + session.unLock(); + } + } + } + + @Override + public Map pullResults(String sessionId, String consumerId) { + if (StringUtils.isBlank(consumerId)) { + return createErrorResult(null, "Consumer ID is null or empty"); + } + + try { + Session session = getCurrentSession(sessionId, false); + SharingResultDistributor resultDistributor = session.getResultDistributor(); + if (resultDistributor == null) { + return createErrorResult(null, "No result distributor found for session: " + sessionId); + } + + ResultConsumer consumer = resultDistributor.getConsumer(consumerId); + if (consumer == null) { + return createErrorResult(null, "Consumer not found: " + consumerId); + } + + List results = consumer.pollResults(); + + if (results != null && results.isEmpty()) { + logger.debug("Filtered empty result list for session: {}, consumer: {}", sessionId, consumerId); + return null; + } + + Map result = new TreeMap<>(); + result.put("success", true); + result.put("sessionId", sessionId); + result.put("consumerId", consumerId); + result.put("results", results); + + Job foregroundJob = session.getForegroundJob(); + if (foregroundJob != null) { + result.put("jobId", foregroundJob.id()); + result.put("jobStatus", foregroundJob.status().toString()); + } + + return result; + + } catch (SessionNotFoundException e) { + return createErrorResult(null, e.getMessage()); + } + } + + @Override + public Map interruptJob(String sessionId) { + try { + Session session = getCurrentSession(sessionId, false); + Job job = session.getForegroundJob(); + if (job == null) { + return createErrorResult(null, "no foreground job is running"); + } + job.interrupt(); + + Map result = new TreeMap<>(); + result.put("success", true); + result.put("sessionId", sessionId); + result.put("jobId", job.id()); + result.put("jobStatus", job.status().toString()); + return result; + + } catch (SessionNotFoundException e) { + return createErrorResult(null, e.getMessage()); + } + } + + @Override + public Map createSession() { + Session session = sessionManager.createSession(); + if (session == null) { + return createErrorResult(null, "create api session failed"); + } + + SharingResultDistributorImpl resultDistributor = new SharingResultDistributorImpl(session); + ResultConsumer resultConsumer = new ResultConsumerImpl(); + resultDistributor.addConsumer(resultConsumer); + session.setResultDistributor(resultDistributor); + + resultDistributor.appendResult(new MessageModel("Welcome to arthas!")); + + WelcomeModel welcomeModel = new WelcomeModel(); + welcomeModel.setVersion(ArthasBanner.version()); + welcomeModel.setWiki(ArthasBanner.wiki()); + welcomeModel.setTutorials(ArthasBanner.tutorials()); + welcomeModel.setMainClass(PidUtils.mainClass()); + welcomeModel.setPid(PidUtils.currentPid()); + welcomeModel.setTime(DateUtils.getCurrentDateTime()); + resultDistributor.appendResult(welcomeModel); + + updateSessionInputStatus(session, InputStatus.ALLOW_INPUT); + + Map result = new TreeMap<>(); + result.put("success", true); + result.put("sessionId", session.getSessionId()); + result.put("consumerId", resultConsumer.getConsumerId()); + return result; + } + + @Override + public Map closeSession(String sessionId) { + try { + Session session = getCurrentSession(sessionId, false); + + if (session.isLocked()) { + session.unLock(); + } + + sessionManager.removeSession(session.getSessionId()); + + Map result = new TreeMap<>(); + result.put("success", true); + result.put("sessionId", sessionId); + return result; + + } catch (SessionNotFoundException e) { + return createErrorResult(null, e.getMessage()); + } + } + private boolean waitForJob(Job job, int timeout) { long startTime = System.currentTimeMillis(); while (true) { @@ -150,16 +326,14 @@ public class CommandExecutorImpl implements CommandExecutor { private Map createErrorResult(String commandLine, String errorMessage) { Map result = new TreeMap<>(); - result.put("command", commandLine); result.put("success", false); result.put("error", errorMessage); - result.put("executionTime", System.currentTimeMillis()); + if (commandLine != null) { + result.put("command", commandLine); + } return result; } - /** - * 创建超时结果 - */ private Map createTimeoutResult(String commandLine, long timeout) { Map result = new TreeMap<>(); result.put("command", commandLine); @@ -170,77 +344,73 @@ public class CommandExecutorImpl implements CommandExecutor { return result; } - private static class JobHandle implements JobListener { - private final CountDownLatch latch = new CountDownLatch(1); - + private void updateSessionInputStatus(Session session, InputStatus inputStatus) { + SharingResultDistributor resultDistributor = session.getResultDistributor(); + if (resultDistributor != null) { + resultDistributor.appendResult(new InputStatusModel(inputStatus)); + } + } + + private Job createJob(String line, Session session, ResultDistributor resultDistributor) { + return createJob(CliTokens.tokenize(line), session, resultDistributor); + } + + private synchronized Job createJob(List args, Session session, ResultDistributor resultDistributor) { + Job job = jobController.createJob(commandManager, args, session, new JobHandler(session), new McpTerm(session), resultDistributor); + return job; + } + + public static class SessionNotFoundException extends RuntimeException { + public SessionNotFoundException(String message) { + super(message); + } + } + + private class JobHandler implements JobListener { + private final Session session; + + public JobHandler(Session session) { + this.session = session; + } + @Override public void onForeground(Job job) { + session.setForegroundJob(job); } @Override public void onBackground(Job job) { + if (session.getForegroundJob() == job) { + session.setForegroundJob(null); + updateSessionInputStatus(session, InputStatus.ALLOW_INPUT); + session.unLock(); + } } @Override public void onTerminated(Job job) { - latch.countDown(); + if (session.getForegroundJob() == job) { + session.setForegroundJob(null); + updateSessionInputStatus(session, InputStatus.ALLOW_INPUT); + session.unLock(); + } } @Override public void onSuspend(Job job) { - } - - public boolean await(long timeout, TimeUnit unit) throws InterruptedException { - return latch.await(timeout, unit); + if (session.getForegroundJob() == job) { + session.setForegroundJob(null); + updateSessionInputStatus(session, InputStatus.ALLOW_INPUT); + session.unLock(); + } } } - public static class InMemoryTerm implements Term { - private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + public static class McpTerm implements Term { private Session session; - private volatile boolean closed = false; - @Override - public Term setSession(Session session) { + public McpTerm(Session session) { this.session = session; - return this; - } - - public Session getSession() { - return session; - } - - @Override - public void readline(String prompt, Handler lineHandler) { - } - - @Override - public void readline(String prompt, Handler lineHandler, Handler completionHandler) { - } - - @Override - public Term closeHandler(Handler 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 @@ -248,29 +418,48 @@ public class CommandExecutorImpl implements CommandExecutor { return this; } + @Override + public String type() { + return "mcp"; + } + + @Override + public int width() { + return 1000; + } + + @Override + public int height() { + return 200; + } + @Override public Term stdinHandler(Handler handler) { return this; } @Override - public Term stdoutHandler(io.termd.core.function.Function handler) { + public Term stdoutHandler(Function 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); - } + public Term write(String data) { + return this; + } + + @Override + public long lastAccessedTime() { + return session.getLastAccessTime(); + } + + @Override + public Term echo(String text) { + return this; + } + + @Override + public Term setSession(Session session) { return this; } @@ -285,35 +474,23 @@ public class CommandExecutorImpl implements CommandExecutor { } @Override - public synchronized void close() { - closed = true; - try { - outputStream.close(); - } catch (Exception e) { - logger.error("Error closing output stream", e); - } + public void readline(String prompt, Handler lineHandler) { + } @Override - public Term echo(String text) { - return write(text); + public void readline(String prompt, Handler lineHandler, Handler completionHandler) { + } - public synchronized String getOutput() { - try { - return outputStream.toString(StandardCharsets.UTF_8.name()); - } catch (UnsupportedEncodingException e) { - logger.error("Error getting output", e); - return ""; - } + @Override + public Term closeHandler(Handler handler) { + return this; } - public synchronized void clearOutput() { - outputStream.reset(); - } - - public boolean isClosed() { - return closed; + @Override + public void close() { + } } } diff --git a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java index 8ee491a74..4a49deb3a 100644 --- a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java +++ b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java @@ -80,8 +80,7 @@ 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 com.taobao.arthas.mcp.server.protocol.server.handler.McpHttpRequestHandler; import io.netty.channel.ChannelFuture; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.util.concurrent.DefaultThreadFactory; @@ -130,7 +129,7 @@ public class ArthasBootstrap { private HttpApiHandler httpApiHandler; - private McpRequestHandler mcpRequestHandler; + private McpHttpRequestHandler mcpRequestHandler; private HttpSessionManager httpSessionManager; private SecurityAuthenticator securityAuthenticator; @@ -683,7 +682,7 @@ public class ArthasBootstrap { return httpApiHandler; } - public McpRequestHandler getMcpRequestHandler() { + public McpHttpRequestHandler getMcpRequestHandler() { return mcpRequestHandler; } diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java index 84c27cac7..043232ecb 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java @@ -5,7 +5,7 @@ import com.alibaba.arthas.deps.org.slf4j.LoggerFactory; import com.taobao.arthas.common.IOUtils; import com.taobao.arthas.core.server.ArthasBootstrap; import com.taobao.arthas.core.shell.term.impl.http.api.HttpApiHandler; -import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpHttpRequestHandler; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelFutureListener; import io.netty.channel.ChannelHandlerContext; @@ -38,7 +38,7 @@ public class HttpRequestHandler extends SimpleChannelInboundHandlerjackson-databind 2.18.1 - - + + com.fasterxml.jackson.datatype + jackson-datatype-jsr310 + 2.18.1 + diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpBootstrap.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpBootstrap.java index dc45ab6a7..00656766c 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpBootstrap.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpBootstrap.java @@ -29,11 +29,18 @@ public class ArthasMcpBootstrap { } public ArthasMcpServer start() { + logger.info("Initializing Arthas MCP Bootstrap..."); try { - // Create and start MCP server - mcpServer = new ArthasMcpServer(); + logger.debug("Creating MCP server instance with command executor: {}", + commandExecutor.getClass().getSimpleName()); + + // Create and start MCP server with CommandExecutor + mcpServer = new ArthasMcpServer(commandExecutor); + logger.debug("MCP server instance created successfully"); + mcpServer.start(); logger.info("Arthas MCP server initialized successfully"); + logger.info("Bootstrap ready - server is operational"); return mcpServer; } catch (Exception e) { logger.error("Failed to initialize Arthas MCP server", e); @@ -42,8 +49,14 @@ public class ArthasMcpBootstrap { } public void shutdown() { + logger.info("Initiating Arthas MCP Bootstrap shutdown..."); if (mcpServer != null) { + logger.debug("Stopping MCP server..."); mcpServer.stop(); + logger.info("MCP server stopped"); + } else { + logger.warn("MCP server was null during shutdown - may not have been properly initialized"); } + logger.info("Arthas MCP Bootstrap shutdown completed"); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpServer.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpServer.java index 212f9aaf5..0593df2f7 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpServer.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/ArthasMcpServer.java @@ -4,20 +4,24 @@ 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.protocol.server.McpStatelessNettyServer; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpHttpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpStatelessHttpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpStreamableHttpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.transport.NettyStatelessServerTransport; +import com.taobao.arthas.mcp.server.protocol.server.transport.NettyStreamableServerTransportProvider; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.Implementation; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.ServerCapabilities; +import com.taobao.arthas.mcp.server.protocol.spec.McpStreamableServerTransportProvider; 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 com.taobao.arthas.mcp.server.util.JsonParser; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Arrays; -import java.util.List; -import java.util.Objects; +import java.util.*; import java.util.stream.Collectors; /** @@ -26,54 +30,52 @@ import java.util.stream.Collectors; */ public class ArthasMcpServer { private static final Logger logger = LoggerFactory.getLogger(ArthasMcpServer.class); - private McpNettyServer server; + + private McpNettyServer streamableServer; + private McpStatelessNettyServer statelessServer; + private final int port; private final String bindAddress; - private McpRequestHandler mcpRequestHandler; + + private final CommandExecutor commandExecutor; + + private McpHttpRequestHandler unifiedMcpHandler; + + private McpStreamableHttpRequestHandler streamableHandler; + + private McpStatelessHttpRequestHandler statelessHandler; - public ArthasMcpServer() { - this(8080, "localhost"); + public ArthasMcpServer(CommandExecutor commandExecutor) { + this(8080, "localhost", commandExecutor); } - public ArthasMcpServer(int port, String bindAddress) { + public ArthasMcpServer(int port, String bindAddress, CommandExecutor commandExecutor) { this.port = port; this.bindAddress = bindAddress; + this.commandExecutor = commandExecutor; } - public McpRequestHandler getMcpRequestHandler() { - return mcpRequestHandler; + public McpHttpRequestHandler getMcpRequestHandler() { + return unifiedMcpHandler; } /** * Start MCP server */ public void start() { + logger.info("Starting Arthas MCP server on {}:{}", bindAddress, port); 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") + .mcpEndpoint("/mcp") .toolChangeNotification(true) .resourceChangeNotification(true) .promptChangeNotification(true) - .objectMapper(new ObjectMapper()) + .objectMapper(JsonParser.getObjectMapper()) .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(); @@ -81,11 +83,53 @@ public class ArthasMcpServer { .filter(Objects::nonNull) .collect(Collectors.toList()); - serverBuilder.tools( - McpToolUtils.toToolSpecifications(providerToolCallbacks, properties)); - server = serverBuilder.build(); + // Create transport for both streamable and stateless servers + McpStreamableServerTransportProvider transportProvider = createStreamableHttpTransportProvider(properties); + streamableHandler = transportProvider.getMcpRequestHandler(); + + NettyStatelessServerTransport statelessTransport = createStatelessHttpTransport(properties); + statelessHandler = statelessTransport.getMcpRequestHandler(); + + unifiedMcpHandler = McpHttpRequestHandler.builder() + .mcpEndpoint(properties.getMcpEndpoint()) + .objectMapper(properties.getObjectMapper()) + .tools(Arrays.asList(callbacks)) + .build(); + unifiedMcpHandler.setStreamableHandler(streamableHandler); + unifiedMcpHandler.setStatelessHandler(statelessHandler); + + // Set up unified MCP handler for both streamable and stateless servers + McpServer.StreamableServerNettySpecification streamableServerNettySpecification = McpServer.netty(transportProvider) + .serverInfo(new Implementation(properties.getName(), properties.getVersion())) + .capabilities(buildServerCapabilities(properties)) + .instructions(properties.getInstructions()) + .requestTimeout(properties.getRequestTimeout()) + .commandExecutor(commandExecutor) + .objectMapper(properties.getObjectMapper() != null ? properties.getObjectMapper() : JsonParser.getObjectMapper()); + + // Set up unified MCP handler for both streamable and stateless servers + McpServer.StatelessServerNettySpecification statelessServerNettySpecification = McpServer.netty(statelessTransport) + .serverInfo(new Implementation(properties.getName(), properties.getVersion())) + .capabilities(buildServerCapabilities(properties)) + .instructions(properties.getInstructions()) + .requestTimeout(properties.getRequestTimeout()) + .commandExecutor(commandExecutor) + .objectMapper(properties.getObjectMapper() != null ? properties.getObjectMapper() : JsonParser.getObjectMapper()); + + streamableServerNettySpecification.tools( + McpToolUtils.toStreamableToolSpecifications(providerToolCallbacks)); + statelessServerNettySpecification.tools( + McpToolUtils.toStatelessToolSpecifications(providerToolCallbacks)); + + streamableServer = streamableServerNettySpecification.build(); + statelessServer = statelessServerNettySpecification.build(); - logger.info("Arthas MCP server started, listening on {}:{}", bindAddress, port); + logger.info("Arthas MCP server started successfully"); + logger.info("- Listening on {}:{}", bindAddress, port); + logger.info("- MCP Endpoint: {}", properties.getMcpEndpoint()); + logger.info("- Transport modes: Streamable + Stateless"); + logger.info("- Available tools: {}", providerToolCallbacks.size()); + logger.info("- Server ready to accept connections"); } catch (Exception e) { logger.error("Failed to start Arthas MCP server", e); throw new RuntimeException("Failed to start Arthas MCP server", e); @@ -95,17 +139,20 @@ public class ArthasMcpServer { /** * Create HTTP transport provider */ - private HttpNettyServerTransportProvider createHttpTransportProvider(McpServerProperties properties) { - return HttpNettyServerTransportProvider.builder() - .messageEndpoint(properties.getMessageEndpoint()) - .sseEndpoint(properties.getSseEndpoint()) + private NettyStreamableServerTransportProvider createStreamableHttpTransportProvider(McpServerProperties properties) { + return NettyStreamableServerTransportProvider.builder() + .mcpEndpoint(properties.getMcpEndpoint()) + .objectMapper(properties.getObjectMapper() != null ? properties.getObjectMapper() : new ObjectMapper()) + .build(); + } + + private NettyStatelessServerTransport createStatelessHttpTransport(McpServerProperties properties) { + return NettyStatelessServerTransport.builder() + .mcpEndpoint(properties.getMcpEndpoint()) .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())) @@ -113,29 +160,31 @@ public class ArthasMcpServer { .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 + public void stop() { + logger.info("Stopping Arthas MCP server..."); try { - Thread.sleep(Long.MAX_VALUE); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); + if (unifiedMcpHandler != null) { + logger.debug("Shutting down unified MCP handler"); + unifiedMcpHandler.closeGracefully().get(); + logger.info("Unified MCP handler stopped successfully"); + } + + if (streamableServer != null) { + logger.debug("Shutting down streamable server"); + streamableServer.closeGracefully().get(); + logger.info("Streamable server stopped successfully"); + } + + if (statelessServer != null) { + logger.debug("Shutting down stateless server"); + statelessServer.closeGracefully().get(); + logger.info("Stateless server stopped successfully"); + } + + logger.info("Arthas MCP server stopped completely"); + } catch (Exception e) { + logger.error("Failed to stop Arthas MCP server gracefully", e); } } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java index c36c2d483..2e35e5e55 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java @@ -2,8 +2,27 @@ package com.taobao.arthas.mcp.server; import java.util.Map; - +/** + * 命令执行器接口 + * + * @author Yeaury 2025/5/26 + */ public interface CommandExecutor { - Map execute(String commandLine, long timeout); + default Map executeSync(String commandLine, long timeout) { + return executeSync(commandLine, timeout, null); + } + + Map executeSync(String commandLine, long timeout, String sessionId); + + Map executeAsync(String commandLine, String sessionId); + + Map pullResults(String sessionId, String consumerId); + + Map interruptJob(String sessionId); + + Map createSession(); + + Map closeSession(String sessionId); } + diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/config/McpServerProperties.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/config/McpServerProperties.java index 94b856dd2..d1bc0d8fc 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/config/McpServerProperties.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/config/McpServerProperties.java @@ -1,7 +1,6 @@ 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; @@ -35,8 +34,7 @@ public class McpServerProperties { */ private final String bindAddress; private final int port; - private final String messageEndpoint; - private final String sseEndpoint; + private final String mcpEndpoint; /** * Timeout configuration @@ -46,6 +44,11 @@ public class McpServerProperties { private final ObjectMapper objectMapper; + /** + * (Optional) response MIME type per tool name. + */ + private Map toolResponseMimeType = new HashMap<>(); + /** * Private constructor, can only be created through Builder */ @@ -59,8 +62,7 @@ public class McpServerProperties { this.resourceSubscribe = builder.resourceSubscribe; this.bindAddress = builder.bindAddress; this.port = builder.port; - this.messageEndpoint = builder.messageEndpoint; - this.sseEndpoint = builder.sseEndpoint; + this.mcpEndpoint = builder.mcpEndpoint; this.requestTimeout = builder.requestTimeout; this.initializationTimeout = builder.initializationTimeout; this.objectMapper = builder.objectMapper; @@ -145,20 +147,13 @@ public class McpServerProperties { return port; } - /** - * Get message endpoint - * @return Message endpoint - */ - public String getMessageEndpoint() { - return messageEndpoint; - } /** * Get SSE endpoint * @return SSE endpoint */ - public String getSseEndpoint() { - return sseEndpoint; + public String getMcpEndpoint() { + return mcpEndpoint; } /** @@ -185,6 +180,14 @@ public class McpServerProperties { return objectMapper; } + public Map getToolResponseMimeType() { + return toolResponseMimeType; + } + + public void setToolResponseMimeType(Map toolResponseMimeType) { + this.toolResponseMimeType = toolResponseMimeType; + } + /** * Builder class for McpServerProperties */ @@ -194,13 +197,12 @@ public class McpServerProperties { private String version = "1.0.0"; private String instructions; private boolean toolChangeNotification = true; - private boolean resourceChangeNotification = true; - private boolean promptChangeNotification = true; + private boolean resourceChangeNotification = false; + private boolean promptChangeNotification = false; private boolean resourceSubscribe = false; private String bindAddress = "localhost"; private int port = 8080; - private String messageEndpoint = "/mcp/message"; - private String sseEndpoint = "/mcp/sse"; + private String mcpEndpoint = "/mcp"; private Duration requestTimeout = Duration.ofSeconds(10); private Duration initializationTimeout = Duration.ofSeconds(30); private ObjectMapper objectMapper; @@ -254,13 +256,8 @@ public class McpServerProperties { return this; } - public Builder messageEndpoint(String messageEndpoint) { - this.messageEndpoint = messageEndpoint; - return this; - } - - public Builder sseEndpoint(String sseEndpoint) { - this.sseEndpoint = sseEndpoint; + public Builder mcpEndpoint(String mcpEndpoint) { + this.mcpEndpoint = mcpEndpoint; return this; } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpStatelessServerHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpStatelessServerHandler.java new file mode 100644 index 000000000..a6b503053 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpStatelessServerHandler.java @@ -0,0 +1,85 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.CommandExecutor; +import com.taobao.arthas.mcp.server.protocol.spec.McpError; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; + +class DefaultMcpStatelessServerHandler implements McpStatelessServerHandler { + + private static final Logger logger = LoggerFactory.getLogger(DefaultMcpStatelessServerHandler.class); + + Map> requestHandlers; + + Map notificationHandlers; + + private final CommandExecutor commandExecutor; + + public DefaultMcpStatelessServerHandler(Map> requestHandlers, + Map notificationHandlers, + CommandExecutor commandExecutor) { + this.requestHandlers = requestHandlers; + this.notificationHandlers = notificationHandlers; + this.commandExecutor = commandExecutor; + } + + @Override + public CompletableFuture handleRequest(McpTransportContext ctx, McpSchema.JSONRPCRequest req) { + ArthasCommandContext commandContext = createCommandContext(); + McpStatelessRequestHandler handler = requestHandlers.get(req.getMethod()); + if (handler == null) { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally(new McpError("Missing handler for request type: " + req.getMethod())); + return f; + } + try { + @SuppressWarnings("unchecked") + CompletableFuture result = (CompletableFuture) handler + .handle(ctx, commandContext, req.getParams()); + return result.handle((r, ex) -> { + if (ex != null) { + Throwable cause = ex instanceof CompletionException ? ex.getCause() : ex; + return new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, req.getId(), null, + new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, cause.getMessage(), null)); + } + return new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, req.getId(), r, null); + }); + } catch (Throwable t) { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally(t); + return f; + } + } + + @Override + public CompletableFuture handleNotification(McpTransportContext ctx, + McpSchema.JSONRPCNotification note) { + McpStatelessNotificationHandler handler = notificationHandlers.get(note.getMethod()); + if (handler == null) { + logger.warn("Missing handler for notification: {}", note.getMethod()); + return CompletableFuture.completedFuture(null); + } + try { + return handler.handle(ctx, note.getParams()); + } catch (Throwable t) { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally(t); + return f; + } + } + + private ArthasCommandContext createCommandContext() { + return new ArthasCommandContext(commandExecutor); + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpTransportContext.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpTransportContext.java new file mode 100644 index 000000000..d42aef8c5 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/DefaultMcpTransportContext.java @@ -0,0 +1,47 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Default implementation for {@link McpTransportContext} which uses a Thread-safe map. + * Objects of this kind are mutable. + */ +public class DefaultMcpTransportContext implements McpTransportContext { + + private final Map storage; + + /** + * Create an empty instance. + */ + public DefaultMcpTransportContext() { + this.storage = new ConcurrentHashMap<>(); + } + + DefaultMcpTransportContext(Map storage) { + this.storage = storage; + } + + @Override + public Object get(String key) { + return this.storage.get(key); + } + + @Override + public void put(String key, Object value) { + this.storage.put(key, value); + } + + /** + * Allows copying the contents. + * @return new instance with the copy of the underlying map + */ + public McpTransportContext copy() { + return new DefaultMcpTransportContext(new ConcurrentHashMap<>(this.storage)); + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpInitRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpInitRequestHandler.java new file mode 100644 index 000000000..a2c1a4922 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpInitRequestHandler.java @@ -0,0 +1,24 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; + +import java.util.concurrent.CompletableFuture; + +/** + * Handles MCP initialization requests from clients using CompletableFuture for async operations. + * This is the Netty-specific version that doesn't depend on Reactor. + */ +public interface McpInitRequestHandler { + + /** + * Handles the initialization request. + * @param initializeRequest the initialization request by the client + * @return a CompletableFuture that will emit the result of the initialization + */ + CompletableFuture handle(McpSchema.InitializeRequest initializeRequest); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServer.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServer.java index a2c5205bf..675ccbfcc 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServer.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServer.java @@ -1,5 +1,18 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + 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.CommandExecutor; +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 org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.time.Duration; import java.util.*; import java.util.concurrent.CompletableFuture; @@ -8,15 +21,6 @@ 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. * @@ -46,13 +50,12 @@ public class McpNettyServer { private McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.DEBUG; - private List protocolVersions = Collections.singletonList(McpSchema.LATEST_PROTOCOL_VERSION); + private List protocolVersions; - public McpNettyServer( - McpServerTransportProvider mcpTransportProvider, - ObjectMapper objectMapper, - Duration requestTimeout, - McpServerFeatures.McpServerConfig features) { + McpNettyServer(McpStreamableServerTransportProvider mcpTransportProvider, + ObjectMapper objectMapper, Duration requestTimeout, + McpServerFeatures.McpServerConfig features, + CommandExecutor commandExecutor) { this.mcpTransportProvider = mcpTransportProvider; this.objectMapper = objectMapper; this.serverInfo = features.getServerInfo(); @@ -63,13 +66,44 @@ public class McpNettyServer { this.resourceTemplates.addAll(features.getResourceTemplates()); this.prompts.putAll(features.getPrompts()); - Map> requestHandlers = new HashMap<>(); + Map> requestHandlers = prepareRequestHandlers(); + Map notificationHandlers = prepareNotificationHandlers(features); + + this.protocolVersions = Collections.singletonList(mcpTransportProvider.protocolVersion()); + + mcpTransportProvider.setSessionFactory(new DefaultMcpStreamableServerSessionFactory(requestTimeout, + this::initializeRequestHandler, requestHandlers, notificationHandlers, commandExecutor)); + } + + private Map prepareNotificationHandlers(McpServerFeatures.McpServerConfig features) { + Map notificationHandlers = new HashMap<>(); + + notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, + (exchange, commandContext, params) -> CompletableFuture.completedFuture(null)); + + List, CompletableFuture>> 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)); + return notificationHandlers; + } + + private Map> prepareRequestHandlers() { + Map> 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())); + (exchange, commandContext, params) -> CompletableFuture.completedFuture(Collections.emptyMap())); // Add tools API handlers if the tool capability is enabled if (this.serverCapabilities.getTools() != null) { @@ -90,37 +124,15 @@ public class McpNettyServer { requestHandlers.put(McpSchema.METHOD_PROMPT_GET, promptsGetRequestHandler()); } + // Add logging API handlers if the logging capability is enabled if (this.serverCapabilities.getLogging() != null) { requestHandlers.put(McpSchema.METHOD_LOGGING_SET_LEVEL, setLoggerRequestHandler()); } - Map notificationHandlers = new HashMap<>(); - - notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, - (exchange, params) -> CompletableFuture.completedFuture(null)); - - List, CompletableFuture>> 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); - }); + return requestHandlers; } + // --------------------------------------- // Lifecycle Management // --------------------------------------- @@ -165,9 +177,9 @@ public class McpNettyServer { this.mcpTransportProvider.close(); } - private McpServerSession.NotificationHandler rootsListChangedNotificationHandler( + private McpNotificationHandler rootsListChangedNotificationHandler( List, CompletableFuture>> rootsChangeConsumers) { - return (exchange, params) -> { + return (exchange, commandContext, params) -> { CompletableFuture futureRoots = exchange.listRoots(); return futureRoots.thenCompose(listRootsResult -> { @@ -268,8 +280,8 @@ public class McpNettyServer { return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED, null); } - private McpServerSession.RequestHandler toolsListRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler toolsListRequestHandler() { + return (exchange, commandContext, params) -> { List tools = new ArrayList<>(); for (McpServerFeatures.ToolSpecification toolSpec : this.tools) { tools.add(toolSpec.getTool()); @@ -279,8 +291,8 @@ public class McpNettyServer { }; } - private McpServerSession.RequestHandler toolsCallRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler toolsCallRequestHandler() { + return (exchange, commandContext, params) -> { McpSchema.CallToolRequest callToolRequest = objectMapper.convertValue(params, new TypeReference() { }); @@ -294,7 +306,7 @@ public class McpNettyServer { return future; } - return toolSpecification.get().getCall().apply(exchange, callToolRequest.getArguments()); + return toolSpecification.get().getCall().apply(exchange, commandContext, callToolRequest); }; } @@ -369,8 +381,8 @@ public class McpNettyServer { return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null); } - private McpServerSession.RequestHandler resourcesListRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler resourcesListRequestHandler() { + return (exchange, commandContext, params) -> { List resourceList = new ArrayList<>(); for (McpServerFeatures.ResourceSpecification spec : this.resources.values()) { resourceList.add(spec.getResource()); @@ -379,13 +391,13 @@ public class McpNettyServer { }; } - private McpServerSession.RequestHandler resourceTemplateListRequestHandler() { - return (exchange, params) -> CompletableFuture + private McpRequestHandler resourceTemplateListRequestHandler() { + return (exchange, commandContext, params) -> CompletableFuture .completedFuture(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null)); } - private McpServerSession.RequestHandler resourcesReadRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler resourcesReadRequestHandler() { + return (exchange, commandContext, params) -> { McpSchema.ReadResourceRequest resourceRequest = objectMapper.convertValue(params, new TypeReference() { }); @@ -473,8 +485,8 @@ public class McpNettyServer { return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_PROMPTS_LIST_CHANGED, null); } - private McpServerSession.RequestHandler promptsListRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler promptsListRequestHandler() { + return (exchange, commandContext, params) -> { List promptList = new ArrayList<>(); for (McpServerFeatures.PromptSpecification promptSpec : this.prompts.values()) { promptList.add(promptSpec.getPrompt()); @@ -483,8 +495,8 @@ public class McpNettyServer { }; } - private McpServerSession.RequestHandler promptsGetRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler promptsGetRequestHandler() { + return (exchange, commandContext, params) -> { McpSchema.GetPromptRequest promptRequest = objectMapper.convertValue(params, new TypeReference() { }); @@ -514,8 +526,8 @@ public class McpNettyServer { loggingMessageNotification); } - private McpServerSession.RequestHandler> setLoggerRequestHandler() { - return (exchange, params) -> { + private McpRequestHandler> setLoggerRequestHandler() { + return (exchange, commandContext, params) -> { try { McpSchema.SetLevelRequest request = this.objectMapper.convertValue(params, McpSchema.SetLevelRequest.class); diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServerExchange.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServerExchange.java index 45a31035b..50f5fd0fc 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServerExchange.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNettyServerExchange.java @@ -1,25 +1,26 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + 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.protocol.spec.McpSession; import com.taobao.arthas.mcp.server.util.Assert; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Map; 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. + * This class is focused only on MCP protocol communication and does not handle command execution directly. * - * @author Yeaury *

* McpNettyServerExchange provides various methods for communicating with the client, including: *

    @@ -28,6 +29,7 @@ import java.util.concurrent.atomic.AtomicLong; *
  • Handling logging notifications *
  • Creating client messages *
  • Managing root directories + *
  • Streamable task management *
*

* Each exchange object is associated with a specific client session, providing context and capabilities for that session. @@ -36,12 +38,16 @@ public class McpNettyServerExchange { private static final Logger logger = LoggerFactory.getLogger(McpNettyServerExchange.class); - private final McpServerSession session; + private final String sessionId; + + private final McpSession session; private final McpSchema.ClientCapabilities clientCapabilities; private final McpSchema.Implementation clientInfo; + private final McpTransportContext transportContext; + private volatile LoggingLevel minLoggingLevel = LoggingLevel.INFO; private static final TypeReference CREATE_MESSAGE_RESULT_TYPE_REF = @@ -52,35 +58,18 @@ public class McpNettyServerExchange { new TypeReference() { }; - /** - * 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"); + public static final TypeReference OBJECT_TYPE_REF = new TypeReference() { + }; + + public McpNettyServerExchange(String sessionId, McpSession session, + McpSchema.ClientCapabilities clientCapabilities, McpSchema.Implementation clientInfo, + McpTransportContext transportContext) { + this.sessionId = sessionId; this.session = session; this.clientCapabilities = clientCapabilities; this.clientInfo = clientInfo; - logger.debug("Created new server exchange, session ID: {}, client: {}", session.getId(), clientInfo); + this.transportContext = transportContext; } - - /** - * 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 @@ -98,23 +87,11 @@ public class McpNettyServerExchange { } /** - * 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 + * Get the MCP server session associated with this exchange. + * @return The MCP server session */ - public CompletableFuture 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()); - } - }); + public McpSession getSession() { + return this.session; } /** @@ -141,15 +118,15 @@ public class McpNettyServerExchange { return future; } - logger.debug("Creating client message, session ID: {}", this.session.getId()); + logger.debug("Creating client message, session ID: {}", this.sessionId); 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()); + logger.error("Failed to create message, session ID: {}, error: {}", this.sessionId, error.getMessage()); } else { - logger.debug("Message created successfully, session ID: {}", this.session.getId()); + logger.debug("Message created successfully, session ID: {}", this.sessionId); } }); } @@ -168,16 +145,16 @@ public class McpNettyServerExchange { * @return Emits a CompletableFuture containing the results of the root list */ public CompletableFuture listRoots(String cursor) { - logger.debug("Requesting root list, session ID: {}, cursor: {}", this.session.getId(), cursor); + logger.debug("Requesting root list, session ID: {}, cursor: {}", this.sessionId, 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()); + logger.error("Failed to get root list, session ID: {}, error: {}", this.sessionId, error.getMessage()); } else { - logger.debug("Root list retrieved successfully, session ID: {}", this.session.getId()); + logger.debug("Root list retrieved successfully, session ID: {}", this.sessionId); } }); } @@ -194,16 +171,20 @@ public class McpNettyServerExchange { .whenComplete((result, error) -> { if (error != null) { logger.error("Failed to send logging notification, level: {}, session ID: {}, error: {}", loggingMessageNotification.getLevel(), - this.session.getId(), error.getMessage()); + this.sessionId, error.getMessage()); } }); } return CompletableFuture.completedFuture(null); } + public CompletableFuture ping() { + return this.session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF); + } + public void setMinLoggingLevel(LoggingLevel minLoggingLevel) { Assert.notNull(minLoggingLevel, "最低日志级别不能为空"); - logger.debug("Setting minimum logging level: {}, session ID: {}", minLoggingLevel, this.session.getId()); + logger.debug("Setting minimum logging level: {}, session ID: {}", minLoggingLevel, this.sessionId); this.minLoggingLevel = minLoggingLevel; } @@ -211,4 +192,21 @@ public class McpNettyServerExchange { return loggingLevel.level() >= this.minLoggingLevel.level(); } + public CompletableFuture progressNotification(McpSchema.ProgressNotification progressNotification) { + if (progressNotification == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("进度通知不能为空")); + return future; + } + + return this.session + .sendNotification(McpSchema.METHOD_NOTIFICATION_PROGRESS, progressNotification) + .whenComplete((result, error) -> { + if (error != null) { + logger.error("Failed to send progress notification, session ID: {}, error: {}", this.sessionId, error.getMessage()); + } else { + logger.debug("Progress notification sent successfully, session ID: {}", this.sessionId); + } + }); + } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNotificationHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNotificationHandler.java new file mode 100644 index 000000000..370a44b4d --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpNotificationHandler.java @@ -0,0 +1,26 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; + +import java.util.concurrent.CompletableFuture; + +/** + * Handles MCP notifications from clients using CompletableFuture for async operations. + * This is the Netty-specific version that doesn't depend on Reactor. + */ +public interface McpNotificationHandler { + + /** + * Handles a notification from the client. + * @param exchange the exchange associated with the client that allows calling back to + * the connected client or inspecting its capabilities. + * @param params the parameters of the notification. + * @return a CompletableFuture that completes once the notification is handled. + */ + CompletableFuture handle(McpNettyServerExchange exchange, ArthasCommandContext arthasCommandContext, Object params); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpRequestHandler.java new file mode 100644 index 000000000..cf1a18b42 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpRequestHandler.java @@ -0,0 +1,26 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; + +import java.util.concurrent.CompletableFuture; + +/** + * Handles MCP requests from clients using CompletableFuture for async operations. + * This is the Netty-specific version that doesn't depend on Reactor. + */ +public interface McpRequestHandler { + + /** + * Handles a request from the client. + * @param exchange the exchange associated with the client that allows calling back to + * the connected client or inspecting its capabilities. + * @param params the parameters of the request. + * @return a CompletableFuture that will emit the response to the request. + */ + CompletableFuture handle(McpNettyServerExchange exchange, ArthasCommandContext arthasCommandContext, Object params); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServer.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServer.java index 10dc77374..bc01aa32d 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServer.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServer.java @@ -1,7 +1,15 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + 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.JsonParser; +import com.taobao.arthas.mcp.server.CommandExecutor; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.protocol.spec.McpStatelessServerTransport; +import com.taobao.arthas.mcp.server.protocol.spec.McpStreamableServerTransportProvider; import com.taobao.arthas.mcp.server.util.Assert; import java.time.Duration; @@ -16,106 +24,77 @@ import java.util.function.BiFunction; */ public interface McpServer { - static NettySpecification netty(McpServerTransportProvider transportProvider) { - return new NettySpecification(transportProvider); + McpSchema.Implementation DEFAULT_SERVER_INFO = new McpSchema.Implementation("mcp-server", "1.0.0"); + + static StreamableServerNettySpecification netty(McpStreamableServerTransportProvider transportProvider) { + return new StreamableServerNettySpecification(transportProvider); } - /** - * serverSpecification - */ - class NettySpecification { + static StatelessServerNettySpecification netty(McpStatelessServerTransport transport) { + return new StatelessServerNettySpecification(transport); + } - private static final McpSchema.Implementation DEFAULT_SERVER_INFO = new McpSchema.Implementation("mcp-server", - "1.0.0"); + class StreamableServerNettySpecification { - private final McpServerTransportProvider transportProvider; + ObjectMapper objectMapper; - private ObjectMapper objectMapper; + McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO; - private McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO; + McpSchema.ServerCapabilities serverCapabilities; - private McpSchema.ServerCapabilities serverCapabilities; + String instructions; - private String instructions; + CommandExecutor commandExecutor; - /** - * 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 tools = new ArrayList<>(); + private final McpStreamableServerTransportProvider transportProvider; - /** - * 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 resources = new HashMap<>(); + final List tools = new ArrayList<>(); - private final List resourceTemplates = new ArrayList<>(); + final Map resources = new HashMap<>(); - /** - * 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 prompts = new HashMap<>(); + final List resourceTemplates = new ArrayList<>(); - private final List, CompletableFuture>> rootsChangeHandlers = new ArrayList<>(); + final Map prompts = new HashMap<>(); - private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout + final List, CompletableFuture>> rootsChangeHandlers = new ArrayList<>(); - private NettySpecification(McpServerTransportProvider transportProvider) { - Assert.notNull(transportProvider, "Transport provider must not be null"); + Duration requestTimeout = Duration.ofSeconds(10); // Default timeout + + public StreamableServerNettySpecification(McpStreamableServerTransportProvider transportProvider) { 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) { + public StreamableServerNettySpecification serverInfo(McpSchema.Implementation serverInfo) { Assert.notNull(serverInfo, "Server info must not be null"); this.serverInfo = serverInfo; return this; } - public NettySpecification requestTimeout(Duration requestTimeout) { + public StreamableServerNettySpecification requestTimeout(Duration requestTimeout) { Assert.notNull(requestTimeout, "Request timeout must not be null"); this.requestTimeout = requestTimeout; return this; } - public NettySpecification serverInfo(String name, String version) { + public StreamableServerNettySpecification 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) { + public StreamableServerNettySpecification instructions(String instructions) { this.instructions = instructions; return this; } - public NettySpecification capabilities(McpSchema.ServerCapabilities serverCapabilities) { + public StreamableServerNettySpecification capabilities(McpSchema.ServerCapabilities serverCapabilities) { this.serverCapabilities = serverCapabilities; return this; } - public NettySpecification tool(McpSchema.Tool tool, - BiFunction, CompletableFuture> handler) { + public StreamableServerNettySpecification tool(McpSchema.Tool tool, + McpServerFeatures.ToolCallFunction handler) { Assert.notNull(tool, "Tool must not be null"); Assert.notNull(handler, "Handler must not be null"); @@ -124,26 +103,24 @@ public interface McpServer { return this; } - public NettySpecification tools(List toolRegistrations) { + public StreamableServerNettySpecification tools(List 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); - } + public StreamableServerNettySpecification tools(McpServerFeatures.ToolSpecification... toolRegistrations) { + this.tools.addAll(Arrays.asList(toolRegistrations)); return this; } - public NettySpecification resources(Map resourceSpecifications) { + public StreamableServerNettySpecification resources(Map resourceSpecifications) { Assert.notNull(resourceSpecifications, "Resource handlers map must not be null"); this.resources.putAll(resourceSpecifications); return this; } - public NettySpecification resources(List resourceSpecifications) { + public StreamableServerNettySpecification resources(List resourceSpecifications) { Assert.notNull(resourceSpecifications, "Resource handlers list must not be null"); for (McpServerFeatures.ResourceSpecification resource : resourceSpecifications) { this.resources.put(resource.getResource().getUri(), resource); @@ -151,7 +128,7 @@ public interface McpServer { return this; } - public NettySpecification resources(McpServerFeatures.ResourceSpecification... resourceSpecifications) { + public StreamableServerNettySpecification 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); @@ -159,25 +136,25 @@ public interface McpServer { return this; } - public NettySpecification resourceTemplates(List resourceTemplates) { + public StreamableServerNettySpecification resourceTemplates(List resourceTemplates) { Assert.notNull(resourceTemplates, "Resource templates must not be null"); this.resourceTemplates.addAll(resourceTemplates); return this; } - public NettySpecification resourceTemplates(McpSchema.ResourceTemplate... resourceTemplates) { + public StreamableServerNettySpecification 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 prompts) { + public StreamableServerNettySpecification prompts(Map prompts) { Assert.notNull(prompts, "Prompts map must not be null"); this.prompts.putAll(prompts); return this; } - public NettySpecification prompts(List prompts) { + public StreamableServerNettySpecification prompts(List prompts) { Assert.notNull(prompts, "Prompts map must not be null"); for (McpServerFeatures.PromptSpecification prompt : prompts) { this.prompts.put(prompt.getPrompt().getName(), prompt); @@ -185,7 +162,7 @@ public interface McpServer { return this; } - public NettySpecification prompts(McpServerFeatures.PromptSpecification... prompts) { + public StreamableServerNettySpecification 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); @@ -193,7 +170,7 @@ public interface McpServer { return this; } - public NettySpecification rootsChangeHandler( + public StreamableServerNettySpecification rootsChangeHandler( BiFunction, CompletableFuture> handler) { Assert.notNull(handler, "Consumer must not be null"); this.rootsChangeHandlers.add(handler); @@ -201,42 +178,222 @@ public interface McpServer { } - public NettySpecification rootsChangeHandlers( + public StreamableServerNettySpecification rootsChangeHandlers( List, CompletableFuture>> handlers) { Assert.notNull(handlers, "Handlers list must not be null"); this.rootsChangeHandlers.addAll(handlers); return this; } - public NettySpecification rootsChangeHandlers( + public StreamableServerNettySpecification rootsChangeHandlers( @SuppressWarnings("unchecked") BiFunction, CompletableFuture>... handlers) { Assert.notNull(handlers, "Handlers list must not be null"); return this.rootsChangeHandlers(Arrays.asList(handlers)); } - public NettySpecification objectMapper(ObjectMapper objectMapper) { + public StreamableServerNettySpecification objectMapper(ObjectMapper objectMapper) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); this.objectMapper = objectMapper; return this; } + public StreamableServerNettySpecification commandExecutor(CommandExecutor commandExecutor) { + Assert.notNull(commandExecutor, "CommandExecutor must not be null"); + this.commandExecutor = commandExecutor; + return this; + } + public McpNettyServer build() { - ObjectMapper mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper(); + ObjectMapper mapper = this.objectMapper != null ? this.objectMapper : JsonParser.getObjectMapper(); + Assert.notNull(this.commandExecutor, "CommandExecutor must be set before building"); return new McpNettyServer( - this.transportProvider, + this.transportProvider, mapper, this.requestTimeout, + new McpServerFeatures.McpServerConfig(this.serverInfo, this.serverCapabilities, this.tools, + this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers, this.instructions + ), this.commandExecutor + ); + } + } + + class StatelessServerNettySpecification { + + private final McpStatelessServerTransport transport; + + ObjectMapper objectMapper; + + McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO; + + McpSchema.ServerCapabilities serverCapabilities; + + String instructions; + + CommandExecutor commandExecutor; + + final List tools = new ArrayList<>(); + + final Map resources = new HashMap<>(); + + final List resourceTemplates = new ArrayList<>(); + + final Map prompts = new HashMap<>(); + + final List, CompletableFuture>> rootsChangeHandlers = new ArrayList<>(); + + Duration requestTimeout = Duration.ofSeconds(10); // Default timeout + + StatelessServerNettySpecification(McpStatelessServerTransport transport) { + this.transport = transport; + } + + public StatelessServerNettySpecification serverInfo(McpSchema.Implementation serverInfo) { + Assert.notNull(serverInfo, "Server info must not be null"); + this.serverInfo = serverInfo; + return this; + } + + public StatelessServerNettySpecification requestTimeout(Duration requestTimeout) { + Assert.notNull(requestTimeout, "Request timeout must not be null"); + this.requestTimeout = requestTimeout; + return this; + } + + public StatelessServerNettySpecification 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 StatelessServerNettySpecification instructions(String instructions) { + this.instructions = instructions; + return this; + } + + public StatelessServerNettySpecification capabilities(McpSchema.ServerCapabilities serverCapabilities) { + this.serverCapabilities = serverCapabilities; + return this; + } + + public StatelessServerNettySpecification tools(List toolRegistrations) { + Assert.notNull(toolRegistrations, "Tool handlers list must not be null"); + this.tools.addAll(toolRegistrations); + return this; + } + + public StatelessServerNettySpecification tools(McpStatelessServerFeatures.ToolSpecification... toolRegistrations) { + for (McpStatelessServerFeatures.ToolSpecification tool : toolRegistrations) { + this.tools.add(tool); + } + return this; + } + + public StatelessServerNettySpecification resources(Map resourceSpecifications) { + Assert.notNull(resourceSpecifications, "Resource handlers map must not be null"); + this.resources.putAll(resourceSpecifications); + return this; + } + + public StatelessServerNettySpecification resources(List resourceSpecifications) { + Assert.notNull(resourceSpecifications, "Resource handlers list must not be null"); + for (McpStatelessServerFeatures.ResourceSpecification resource : resourceSpecifications) { + this.resources.put(resource.getResource().getUri(), resource); + } + return this; + } + + public StatelessServerNettySpecification resources(McpStatelessServerFeatures.ResourceSpecification... resourceSpecifications) { + Assert.notNull(resourceSpecifications, "Resource handlers list must not be null"); + for (McpStatelessServerFeatures.ResourceSpecification resource : resourceSpecifications) { + this.resources.put(resource.getResource().getUri(), resource); + } + return this; + } + + public StatelessServerNettySpecification resourceTemplates(List resourceTemplates) { + Assert.notNull(resourceTemplates, "Resource templates must not be null"); + this.resourceTemplates.addAll(resourceTemplates); + return this; + } + + public StatelessServerNettySpecification resourceTemplates(McpSchema.ResourceTemplate... resourceTemplates) { + Assert.notNull(resourceTemplates, "Resource templates must not be null"); + this.resourceTemplates.addAll(Arrays.asList(resourceTemplates)); + return this; + } + + public StatelessServerNettySpecification prompts(Map prompts) { + Assert.notNull(prompts, "Prompts map must not be null"); + this.prompts.putAll(prompts); + return this; + } + + public StatelessServerNettySpecification prompts(List prompts) { + Assert.notNull(prompts, "Prompts map must not be null"); + for (McpStatelessServerFeatures.PromptSpecification prompt : prompts) { + this.prompts.put(prompt.getPrompt().getName(), prompt); + } + return this; + } + + public StatelessServerNettySpecification prompts(McpStatelessServerFeatures.PromptSpecification... prompts) { + Assert.notNull(prompts, "Prompts map must not be null"); + for (McpStatelessServerFeatures.PromptSpecification prompt : prompts) { + this.prompts.put(prompt.getPrompt().getName(), prompt); + } + return this; + } + + public StatelessServerNettySpecification rootsChangeHandler( + BiFunction, CompletableFuture> handler) { + Assert.notNull(handler, "Consumer must not be null"); + this.rootsChangeHandlers.add(handler); + return this; + } + + + public StatelessServerNettySpecification rootsChangeHandlers( + List, CompletableFuture>> handlers) { + Assert.notNull(handlers, "Handlers list must not be null"); + this.rootsChangeHandlers.addAll(handlers); + return this; + } + + public StatelessServerNettySpecification rootsChangeHandlers( + @SuppressWarnings("unchecked") BiFunction, CompletableFuture>... handlers) { + Assert.notNull(handlers, "Handlers list must not be null"); + return this.rootsChangeHandlers(Arrays.asList(handlers)); + } + + public StatelessServerNettySpecification objectMapper(ObjectMapper objectMapper) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + this.objectMapper = objectMapper; + return this; + } + + public StatelessServerNettySpecification commandExecutor(CommandExecutor commandExecutor) { + Assert.notNull(commandExecutor, "CommandExecutor must not be null"); + this.commandExecutor = commandExecutor; + return this; + } + + public McpStatelessNettyServer build() { + ObjectMapper mapper = this.objectMapper != null ? this.objectMapper : JsonParser.getObjectMapper(); + return new McpStatelessNettyServer( + this.transport, mapper, this.requestTimeout, - new McpServerFeatures.McpServerConfig( + new McpStatelessServerFeatures.McpServerConfig( this.serverInfo, this.serverCapabilities, this.tools, this.resources, this.resourceTemplates, this.prompts, - this.rootsChangeHandlers, this.instructions - ) + ), + this.commandExecutor ); } + } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServerFeatures.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServerFeatures.java index efc75014d..9ab7f073c 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServerFeatures.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpServerFeatures.java @@ -1,14 +1,15 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.server; -import com.taobao.arthas.mcp.server.protocol.spec.*; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; 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.*; import java.util.concurrent.CompletableFuture; import java.util.function.BiFunction; @@ -159,11 +160,11 @@ public class McpServerFeatures { public static class ToolSpecification { private final McpSchema.Tool tool; - private final BiFunction, CompletableFuture> call; + private final ToolCallFunction call; public ToolSpecification( McpSchema.Tool tool, - BiFunction, CompletableFuture> call) { + ToolCallFunction call) { this.tool = tool; this.call = call; } @@ -172,11 +173,23 @@ public class McpServerFeatures { return tool; } - public BiFunction, CompletableFuture> getCall() { + public ToolCallFunction getCall() { return call; } } + /** + * Tool call function interface with three parameters + */ + @FunctionalInterface + public interface ToolCallFunction { + CompletableFuture apply( + McpNettyServerExchange exchange, + ArthasCommandContext commandContext, + McpSchema.CallToolRequest arguments + ); + } + public static class ResourceSpecification { private final McpSchema.Resource resource; private final BiFunction> readHandler; diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNettyServer.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNettyServer.java new file mode 100644 index 000000000..b64249c82 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNettyServer.java @@ -0,0 +1,456 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +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.CommandExecutor; +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.McpStatelessServerTransport; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +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; + +/** + * A Netty-based MCP server implementation that provides access to tools, resources, and prompts. + * + * @author Yeaury + */ +public class McpStatelessNettyServer { + + private static final Logger logger = LoggerFactory.getLogger(McpStatelessNettyServer.class); + + private final McpStatelessServerTransport mcpTransportProvider; + + private final ObjectMapper objectMapper; + + private final McpSchema.ServerCapabilities serverCapabilities; + + private final McpSchema.Implementation serverInfo; + + private final String instructions; + + private final CopyOnWriteArrayList tools = new CopyOnWriteArrayList<>(); + + private final CopyOnWriteArrayList resourceTemplates = new CopyOnWriteArrayList<>(); + + private final ConcurrentHashMap resources = new ConcurrentHashMap<>(); + + private final ConcurrentHashMap prompts = new ConcurrentHashMap<>(); + + private List protocolVersions = Collections.singletonList(McpSchema.LATEST_PROTOCOL_VERSION); + + public McpStatelessNettyServer( + McpStatelessServerTransport mcpTransport, + ObjectMapper objectMapper, + Duration requestTimeout, + McpStatelessServerFeatures.McpServerConfig features, + CommandExecutor commandExecutor) { + this.mcpTransportProvider = mcpTransport; + 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> 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, commandContext, 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()); + } + + McpStatelessServerHandler handler = new DefaultMcpStatelessServerHandler(requestHandlers, new HashMap<>(), commandExecutor); + mcpTransport.setMcpHandler(handler); + } + + // --------------------------------------- + // Lifecycle Management + // --------------------------------------- + private CompletableFuture 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 closeGracefully() { + return this.mcpTransportProvider.closeGracefully(); + } + + public void close() { + this.mcpTransportProvider.close(); + } + + private McpNotificationHandler rootsListChangedNotificationHandler( + List, CompletableFuture>> rootsChangeConsumers) { + return (exchange, commandContext, params) -> { + CompletableFuture futureRoots = exchange.listRoots(); + + return futureRoots.thenCompose(listRootsResult -> { + List roots = listRootsResult.getRoots(); + List> futures = new ArrayList<>(); + for (BiFunction, CompletableFuture> consumer : rootsChangeConsumers) { + CompletableFuture 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 addTool(McpStatelessServerFeatures.ToolSpecification toolSpecification) { + if (toolSpecification == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Tool specification must not be null")); + return future; + } + if (toolSpecification.getTool() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Tool must not be null")); + return future; + } + if (toolSpecification.getCall() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Tool call handler must not be null")); + return future; + } + if (this.serverCapabilities.getTools() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Server must be configured with tool capabilities")); + return future; + } + + return CompletableFuture + .runAsync(() -> { + 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()); + }) + .exceptionally(ex -> { + Throwable cause = ex instanceof CompletionException ? ex.getCause() : ex; + logger.error("Error while adding tool", cause); + throw new CompletionException(cause); + }); + } + + + + public CompletableFuture removeTool(String toolName) { + if (toolName == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Tool name must not be null")); + return future; + } + if (this.serverCapabilities.getTools() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Server must be configured with tool capabilities")); + return future; + } + + return CompletableFuture + .runAsync(() -> { + 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); + }) + .exceptionally(ex -> { + Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex; + logger.error("Error while removing tool '{}'", toolName, cause); + throw new CompletionException(cause); + }); + } + + private McpStatelessRequestHandler toolsListRequestHandler() { + return (exchange, commandContext, params) -> { + List tools = new ArrayList<>(); + for (McpStatelessServerFeatures.ToolSpecification toolSpec : this.tools) { + tools.add(toolSpec.getTool()); + } + + return CompletableFuture.completedFuture(new McpSchema.ListToolsResult(tools, null)); + }; + } + + private McpStatelessRequestHandler toolsCallRequestHandler() { + return (context, commandContext, params) -> { + McpSchema.CallToolRequest callToolRequest = objectMapper.convertValue(params, + new TypeReference() { + }); + Optional toolSpecification = this.tools.stream() + .filter(tr -> callToolRequest.getName().equals(tr.getTool().getName())) + .findAny(); + + if (!toolSpecification.isPresent()) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("no tool found: " + callToolRequest.getName())); + return future; + } + + return toolSpecification.get().getCall().apply(context, commandContext, callToolRequest.getArguments()); + }; + } + + // --------------------------------------- + // Resource Management + // --------------------------------------- + + public CompletableFuture addResource(McpStatelessServerFeatures.ResourceSpecification resourceSpecification) { + if (resourceSpecification == null || resourceSpecification.getResource() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Resource must not be null")); + return future; + } + if (this.serverCapabilities.getResources() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Server must be configured with resource capabilities")); + return future; + } + + return CompletableFuture + .runAsync(() -> { + String uri = resourceSpecification.getResource().getUri(); + if (this.resources.putIfAbsent(uri, resourceSpecification) != null) { + throw new CompletionException(new McpError("Resource with URI '" + uri + "' already exists")); + } + logger.debug("Added resource handler: {}", uri); + }) + .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 removeResource(String resourceUri) { + if (resourceUri == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Resource URI must not be null")); + return future; + } + if (this.serverCapabilities.getResources() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Server must be configured with resource capabilities")); + return future; + } + + return CompletableFuture + .runAsync(() -> { + McpStatelessServerFeatures.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); + }) + .exceptionally(ex -> { + Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex; + logger.error("Error while removing resource '{}'", resourceUri, cause); + throw new CompletionException(cause); + }); + } + + private McpStatelessRequestHandler resourcesListRequestHandler() { + return (exchange, commandContext, params) -> { + List resourceList = new ArrayList<>(); + for (McpStatelessServerFeatures.ResourceSpecification spec : this.resources.values()) { + resourceList.add(spec.getResource()); + } + return CompletableFuture.completedFuture(new McpSchema.ListResourcesResult(resourceList, null)); + }; + } + + private McpStatelessRequestHandler resourceTemplateListRequestHandler() { + return (context, commandContext, params) -> CompletableFuture + .completedFuture(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null)); + } + + private McpStatelessRequestHandler resourcesReadRequestHandler() { + return (context, commandContext, params) -> { + McpSchema.ReadResourceRequest resourceRequest = objectMapper.convertValue(params, + new TypeReference() { + }); + String resourceUri = resourceRequest.getUri(); + McpStatelessServerFeatures.ResourceSpecification specification = this.resources.get(resourceUri); + if (specification != null) { + return specification.getReadHandler().apply(context, resourceRequest); + } + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Resource not found: " + resourceUri)); + return future; + }; + } + + // --------------------------------------- + // Prompt Management + // --------------------------------------- + + public CompletableFuture addPrompt(McpStatelessServerFeatures.PromptSpecification promptSpecification) { + if (promptSpecification == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Prompt specification must not be null")); + return future; + } + if (this.serverCapabilities.getPrompts() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Server must be configured with prompt capabilities")); + return future; + } + + return CompletableFuture + .runAsync(() -> { + String name = promptSpecification.getPrompt().getName(); + McpStatelessServerFeatures.PromptSpecification existing = + this.prompts.putIfAbsent(name, promptSpecification); + if (existing != null) { + throw new CompletionException(new McpError("Prompt with name '" + name + "' already exists")); + } + logger.debug("Added prompt handler: {}", name); + }) + .exceptionally(ex -> { + Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex; + String name = promptSpecification.getPrompt().getName(); + logger.error("Error while adding prompt '{}'", name, cause); + throw new CompletionException(cause); + }); + } + + public CompletableFuture removePrompt(String promptName) { + if (promptName == null || promptName.isEmpty()) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Prompt name must not be null or empty")); + return future; + } + if (this.serverCapabilities.getPrompts() == null) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Server must be configured with prompt capabilities")); + return future; + } + + return CompletableFuture + .runAsync(() -> { + McpStatelessServerFeatures.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); + }) + .exceptionally(ex -> { + Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex; + logger.error("Error while removing prompt '{}'", promptName, cause); + throw new CompletionException(cause); + }); + } + + + private McpStatelessRequestHandler promptsListRequestHandler() { + return (exchange, commandContext, params) -> { + List promptList = new ArrayList<>(); + for (McpStatelessServerFeatures.PromptSpecification promptSpec : this.prompts.values()) { + promptList.add(promptSpec.getPrompt()); + } + return CompletableFuture.completedFuture(new McpSchema.ListPromptsResult(promptList, null)); + }; + } + + private McpStatelessRequestHandler promptsGetRequestHandler() { + return (context, commandContext, params) -> { + McpSchema.GetPromptRequest promptRequest = objectMapper.convertValue(params, + new TypeReference() { + }); + + McpStatelessServerFeatures.PromptSpecification specification = this.prompts.get(promptRequest.getName()); + if (specification != null) { + return specification.getPromptHandler().apply(context, promptRequest); + } + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(new McpError("Prompt not found: " + promptRequest.getName())); + return future; + }; + } + + + // --------------------------------------- + // Sampling + // --------------------------------------- + + public void setProtocolVersions(List protocolVersions) { + this.protocolVersions = protocolVersions; + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNotificationHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNotificationHandler.java new file mode 100644 index 000000000..13b21f173 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessNotificationHandler.java @@ -0,0 +1,23 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + + +import java.util.concurrent.CompletableFuture; + +/** + * Handler for MCP notifications in a stateless server. + */ +public interface McpStatelessNotificationHandler { + + /** + * Handle to notification and complete once done. + * @param transportContext {@link McpTransportContext} associated with the transport + * @param params the payload of the MCP notification + * @return Mono which completes once the processing is done + */ + CompletableFuture handle(McpTransportContext transportContext, Object params); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessRequestHandler.java new file mode 100644 index 000000000..bc15d63f3 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessRequestHandler.java @@ -0,0 +1,24 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; + +import java.util.concurrent.CompletableFuture; + +/** + * Handler for MCP requests in a stateless server. + */ +public interface McpStatelessRequestHandler { + + /** + * Handle the request and complete with a result. + * @param transportContext {@link McpTransportContext} associated with the transport + * @param params the payload of the MCP request + * @return Mono which completes with the response object + */ + CompletableFuture handle(McpTransportContext transportContext, ArthasCommandContext arthasCommandContext, Object params); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerFeatures.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerFeatures.java new file mode 100644 index 000000000..e41524a7b --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerFeatures.java @@ -0,0 +1,218 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.util.Assert; +import com.taobao.arthas.mcp.server.util.Utils; + +import java.util.*; +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 McpStatelessServerFeatures { + + public static class McpServerConfig { + private final McpSchema.Implementation serverInfo; + private final McpSchema.ServerCapabilities serverCapabilities; + private final List tools; + private final Map resources; + private final List resourceTemplates; + private final Map prompts; + private final String instructions; + + public McpServerConfig( + McpSchema.Implementation serverInfo, + McpSchema.ServerCapabilities serverCapabilities, + List tools, + Map resources, + List resourceTemplates, + Map prompts, + 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.serverInfo = serverInfo; + this.serverCapabilities = serverCapabilities; + this.instructions = instructions; + } + + public McpSchema.Implementation getServerInfo() { + return serverInfo; + } + + public McpSchema.ServerCapabilities getServerCapabilities() { + return serverCapabilities; + } + + public List getTools() { + return tools; + } + + public Map getResources() { + return resources; + } + + public List getResourceTemplates() { + return resourceTemplates; + } + + public Map getPrompts() { + return prompts; + } + + 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 tools = new ArrayList<>(); + private final Map resources = new HashMap<>(); + private final List resourceTemplates = new ArrayList<>(); + private final Map prompts = new HashMap<>(); + 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 instructions(String instructions) { + this.instructions = instructions; + return this; + } + + public McpServerConfig build() { + return new McpServerConfig(serverInfo, serverCapabilities, tools, resources, resourceTemplates, prompts, + instructions); + } + } + } + + public static class ToolSpecification { + private final McpSchema.Tool tool; + private final ToolCallFunction call; + + public ToolSpecification( + McpSchema.Tool tool, + ToolCallFunction call) { + this.tool = tool; + this.call = call; + } + + public McpSchema.Tool getTool() { + return tool; + } + + public ToolCallFunction getCall() { + return call; + } + } + + /** + * Tool call function interface with three parameters + */ + @FunctionalInterface + public interface ToolCallFunction { + CompletableFuture apply( + McpTransportContext context, + ArthasCommandContext commandContext, + Map arguments + ); + } + + public static class ResourceSpecification { + private final McpSchema.Resource resource; + private final BiFunction> readHandler; + + public ResourceSpecification( + McpSchema.Resource resource, + BiFunction> readHandler) { + this.resource = resource; + this.readHandler = readHandler; + } + + public McpSchema.Resource getResource() { + return resource; + } + + public BiFunction> getReadHandler() { + return readHandler; + } + } + + public static class PromptSpecification { + private final McpSchema.Prompt prompt; + private final BiFunction> promptHandler; + + public PromptSpecification( + McpSchema.Prompt prompt, + BiFunction> promptHandler) { + this.prompt = prompt; + this.promptHandler = promptHandler; + } + + public McpSchema.Prompt getPrompt() { + return prompt; + } + + public BiFunction> getPromptHandler() { + return promptHandler; + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerHandler.java new file mode 100644 index 000000000..821284b9d --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpStatelessServerHandler.java @@ -0,0 +1,32 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; + +import java.util.concurrent.CompletableFuture; + +public interface McpStatelessServerHandler { + + /** + * Handle the request using user-provided feature implementations. + * @param transportContext {@link McpTransportContext} carrying transport layer + * metadata + * @param request the request JSON object + * @return Mono containing the JSON response + */ + CompletableFuture handleRequest(McpTransportContext transportContext, + McpSchema.JSONRPCRequest request); + + /** + * Handle the notification. + * @param transportContext {@link McpTransportContext} carrying transport layer + * metadata + * @param notification the notification JSON object + * @return Mono that completes once handling is finished + */ + CompletableFuture handleNotification(McpTransportContext transportContext, McpSchema.JSONRPCNotification notification); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContext.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContext.java new file mode 100644 index 000000000..a2b36e109 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContext.java @@ -0,0 +1,19 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +public interface McpTransportContext { + + String KEY = "MCP_TRANSPORT_CONTEXT"; + + McpTransportContext EMPTY = new DefaultMcpTransportContext(); + + Object get(String key); + + void put(String key, Object value); + + McpTransportContext copy(); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContextExtractor.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContextExtractor.java new file mode 100644 index 000000000..c9e964c05 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/McpTransportContextExtractor.java @@ -0,0 +1,25 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server; + +/** + * Interface for extracting transport context from server requests. + * This allows inspection of HTTP transport level metadata during request processing. + * + * @param the type of server request + */ +@FunctionalInterface +public interface McpTransportContextExtractor { + + /** + * Extract transport context from the server request. + * + * @param serverRequest the server request to extract context from + * @param context the base context to fill in + * @return the updated context with extracted information + */ + McpTransportContext extract(T serverRequest, McpTransportContext context); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpHttpRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpHttpRequestHandler.java new file mode 100644 index 000000000..b2f135475 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpHttpRequestHandler.java @@ -0,0 +1,324 @@ +package com.taobao.arthas.mcp.server.protocol.server.handler; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContextExtractor; +import com.taobao.arthas.mcp.server.protocol.spec.McpError; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.tool.ToolCallback; +import com.taobao.arthas.mcp.server.util.Assert; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.channel.ChannelHandler; +import io.netty.channel.ChannelHandlerContext; +import io.netty.handler.codec.http.*; +import io.netty.util.CharsetUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; + +import static com.taobao.arthas.mcp.server.protocol.spec.McpSchema.*; + +/** + * MCP HTTP请求处理器，分发请求到无状态或流式处理器。 + * + * @author Yeaury + */ +public class McpHttpRequestHandler { + + private static final Logger logger = LoggerFactory.getLogger(McpHttpRequestHandler.class); + + public static final String APPLICATION_JSON = "application/json"; + public static final String TEXT_EVENT_STREAM = "text/event-stream"; + private static final String ACCEPT_HEADER = "Accept"; + + private final String mcpEndpoint; + private final ObjectMapper objectMapper; + private final McpTransportContextExtractor contextExtractor; + private final AtomicBoolean isClosing = new AtomicBoolean(false); + + private McpStatelessHttpRequestHandler statelessHandler; + + private McpStreamableHttpRequestHandler streamableHandler; + + private final Map toolStreamableCache = new ConcurrentHashMap<>(); + + public McpHttpRequestHandler(String mcpEndpoint, ObjectMapper objectMapper, + McpTransportContextExtractor contextExtractor, + List tools) { + Assert.notNull(mcpEndpoint, "mcpEndpoint must not be null"); + Assert.notNull(objectMapper, "objectMapper must not be null"); + Assert.notNull(contextExtractor, "contextExtractor must not be null"); + + this.mcpEndpoint = mcpEndpoint; + this.objectMapper = objectMapper; + this.contextExtractor = contextExtractor; + registerStreamableTools(tools); + } + + public void setStatelessHandler(McpStatelessHttpRequestHandler statelessHandler) { + this.statelessHandler = statelessHandler; + } + + public void setStreamableHandler(McpStreamableHttpRequestHandler streamableHandler) { + this.streamableHandler = streamableHandler; + } + + public void handle(ChannelHandlerContext ctx, FullHttpRequest request) throws Exception { + String uri = request.uri(); + if (!uri.endsWith(mcpEndpoint)) { + sendError(ctx, HttpResponseStatus.NOT_FOUND, new McpError("Endpoint not found")); + return; + } + + if (isClosing.get()) { + sendError(ctx, HttpResponseStatus.SERVICE_UNAVAILABLE, new McpError("Server is shutting down")); + return; + } + + boolean useStreamable = shouldUseStreamableTransport(request); + + logger.debug("Request {} {} -> using {} transport", + request.method(), request.uri(), useStreamable ? "streamable" : "stateless"); + + try { + if (useStreamable) { + if (streamableHandler == null) { + sendError(ctx, HttpResponseStatus.SERVICE_UNAVAILABLE, + new McpError("Streamable transport handler not available")); + return; + } + streamableHandler.handle(ctx, request); + } else { + if (statelessHandler == null) { + sendError(ctx, HttpResponseStatus.SERVICE_UNAVAILABLE, + new McpError("Stateless transport handler not available")); + return; + } + statelessHandler.handle(ctx, request); + } + } catch (Exception e) { + logger.error("Error handling request: {}", e.getMessage(), e); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Error processing request: " + e.getMessage())); + } + } + + private boolean shouldUseStreamableTransport(FullHttpRequest request) { + String accept = request.headers().get(ACCEPT_HEADER); + boolean supportsSSE = (accept != null && accept.contains(TEXT_EVENT_STREAM)); + HttpMethod method = request.method(); + + if (!supportsSSE) { + logger.debug("Using stateless transport for {} without SSE support", method); + return false; + } + + if (method == HttpMethod.GET || method == HttpMethod.DELETE) { + logger.debug("Using streamable transport for {} with SSE support", method); + return true; + } + + if (method == HttpMethod.POST) { + try { + ByteBuf content = request.content(); + if (content.readableBytes() > 0) { + String body = content.toString(CharsetUtil.UTF_8); + content.resetReaderIndex(); + + McpSchema.JSONRPCMessage message = + McpSchema.deserializeJsonRpcMessage(objectMapper, body); + return isStreamingRequest(message); + } + } catch (Exception e) { + logger.debug("Failed to parse POST body, defaulting to stateless: {}", e.getMessage()); + } + } + + logger.debug("Using stateless transport as default"); + return false; + } + + + /** + * 判断是否为流式工具 + */ + private boolean isStreamingRequest(McpSchema.JSONRPCMessage message) { + if (message instanceof McpSchema.JSONRPCRequest) { + McpSchema.JSONRPCRequest request = (McpSchema.JSONRPCRequest) message; + String method = request.getMethod(); + + if (METHOD_INITIALIZE.equals(method)) { + logger.debug("Initialize method is streamable"); + return true; + } + + if (METHOD_TOOLS_LIST.equals(method)) { + logger.debug("tools/list method is not streamable"); + return false; + } + + if (METHOD_TOOLS_CALL.equals(method)) { + return isStreamableToolCall(request); + } + } + + if (message instanceof McpSchema.JSONRPCNotification || + message instanceof McpSchema.JSONRPCResponse) { + logger.debug("Notification/response message is streamable"); + return true; + } + + return false; + } + + private boolean isStreamableToolCall(McpSchema.JSONRPCRequest request) { + Object params = request.getParams(); + if (!(params instanceof java.util.Map)) { + return false; + } + + @SuppressWarnings("unchecked") + java.util.Map paramMap = (java.util.Map) params; + String toolName = (String) paramMap.get("name"); + + if (toolName == null) { + return false; + } + + boolean isStreamable = isStreamableTool(toolName); + logger.debug("Tool {} is {} (from cache)", toolName, isStreamable ? "streamable" : "non-streamable"); + + return isStreamable; + } + + public void registerStreamableTools(List tools) { + if (tools == null || tools.isEmpty()) { + return; + } + + for (ToolCallback tool : tools) { + if (tool != null && tool.getToolDefinition() != null) { + String toolName = tool.getToolDefinition().getName(); + boolean isStreamable = tool.getToolDefinition().isStreamable(); + toolStreamableCache.put(toolName, isStreamable); + + logger.debug("Registered tool: {} (streamable: {})", toolName, isStreamable); + } + } + + logger.info("Registered {} tools to cache", toolStreamableCache.size()); + } + + private boolean isStreamableTool(String toolName) { + if (toolName == null) { + return false; + } + + Boolean isStreamable = toolStreamableCache.get(toolName); + if (isStreamable != null) { + return isStreamable; + } + + logger.warn("Tool '{}' not found in cache, assuming non-streamable", toolName); + return false; + } + + public CompletableFuture closeGracefully() { + return CompletableFuture.runAsync(() -> { + this.isClosing.set(true); + logger.debug("Initiating graceful shutdown of MCP handler"); + + CompletableFuture statelessClose = CompletableFuture.completedFuture(null); + CompletableFuture streamableClose = CompletableFuture.completedFuture(null); + + if (statelessHandler != null) { + statelessClose = statelessHandler.closeGracefully(); + } + + if (streamableHandler != null) { + streamableClose = streamableHandler.closeGracefully(); + } + + CompletableFuture.allOf(statelessClose, streamableClose).join(); + logger.debug("Graceful shutdown completed"); + }); + } + + private void sendError(ChannelHandlerContext ctx, HttpResponseStatus status, McpError mcpError) { + try { + String jsonError = objectMapper.writeValueAsString(mcpError); + ByteBuf content = Unpooled.copiedBuffer(jsonError, 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, response.content().readableBytes()); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + + ctx.writeAndFlush(response); + } catch (Exception e) { + logger.error("Failed to send error response: {}", e.getMessage()); + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.INTERNAL_SERVER_ERROR + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + ctx.writeAndFlush(response); + } + } + + public String getMcpEndpoint() { + return mcpEndpoint; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private String mcpEndpoint = "/mcp"; + private ObjectMapper objectMapper; + private McpTransportContextExtractor contextExtractor = (request, context) -> context; + private List tools; + + public Builder mcpEndpoint(String mcpEndpoint) { + Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); + this.mcpEndpoint = mcpEndpoint; + return this; + } + + public Builder objectMapper(ObjectMapper objectMapper) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + this.objectMapper = objectMapper; + return this; + } + + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "Context extractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + + public Builder tools(List tools) { + Assert.notNull(tools, "Tools must not be null"); + this.tools = tools; + return this; + } + + public McpHttpRequestHandler build() { + Assert.notNull(this.objectMapper, "ObjectMapper must be set"); + Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set"); + + return new McpHttpRequestHandler(this.mcpEndpoint, this.objectMapper, this.contextExtractor, tools); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpRequestHandler.java deleted file mode 100644 index 9c2d3eebf..000000000 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpRequestHandler.java +++ /dev/null @@ -1,379 +0,0 @@ -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 sessions = new ConcurrentHashMap<>(); - - /** SSE connection mapping: session ID -> Channel */ - private final Map 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 notifyClients(String method, Object params) { - CompletableFuture 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 closeGracefully() { - CompletableFuture 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 sendMessage(McpSchema.JSONRPCMessage message) { - CompletableFuture 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 unmarshalFrom(Object data, TypeReference typeRef) { - return objectMapper.convertValue(data, typeRef); - } - - @Override - public CompletableFuture closeGracefully() { - CompletableFuture 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(); - } - } - - } - -} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStatelessHttpRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStatelessHttpRequestHandler.java new file mode 100644 index 000000000..64246aa25 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStatelessHttpRequestHandler.java @@ -0,0 +1,227 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server.handler; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.taobao.arthas.mcp.server.protocol.server.DefaultMcpTransportContext; +import com.taobao.arthas.mcp.server.protocol.server.McpStatelessServerHandler; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContext; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContextExtractor; +import com.taobao.arthas.mcp.server.protocol.spec.McpError; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.util.Assert; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.channel.ChannelFutureListener; +import io.netty.channel.ChannelHandler; +import io.netty.channel.ChannelHandlerContext; +import io.netty.handler.codec.http.*; +import io.netty.util.CharsetUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Server-side HTTP request handler for stateless MCP transport. + * This handler processes HTTP requests without maintaining client sessions. + * + */ +public class McpStatelessHttpRequestHandler { + + private static final Logger logger = LoggerFactory.getLogger(McpStatelessHttpRequestHandler.class); + + public static final String UTF_8 = "UTF-8"; + public static final String APPLICATION_JSON = "application/json"; + public static final String TEXT_EVENT_STREAM = "text/event-stream"; + public static final String ACCEPT = "Accept"; + private static final String FAILED_TO_SEND_ERROR_RESPONSE = "Failed to send error response: {}"; + + private final ObjectMapper objectMapper; + private final String mcpEndpoint; + private McpStatelessServerHandler mcpHandler; + private final McpTransportContextExtractor contextExtractor; + private final AtomicBoolean isClosing = new AtomicBoolean(false); + + /** + * Constructs a new McpStatelessHttpRequestHandler instance. + * + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * @param mcpEndpoint The endpoint URI where clients should send their JSON-RPC messages + * @param contextExtractor The extractor for transport context from the request + */ + public McpStatelessHttpRequestHandler(ObjectMapper objectMapper, String mcpEndpoint, + McpTransportContextExtractor contextExtractor) { + Assert.notNull(objectMapper, "objectMapper must not be null"); + Assert.notNull(mcpEndpoint, "mcpEndpoint must not be null"); + Assert.notNull(contextExtractor, "contextExtractor must not be null"); + + this.objectMapper = objectMapper; + this.mcpEndpoint = mcpEndpoint; + this.contextExtractor = contextExtractor; + } + + public void setMcpHandler(McpStatelessServerHandler mcpHandler) { + this.mcpHandler = mcpHandler; + } + + /** + * Initiates a graceful shutdown of the handler. + * + * @return A CompletableFuture that completes when shutdown is initiated + */ + public CompletableFuture closeGracefully() { + return CompletableFuture.supplyAsync(() -> { + this.isClosing.set(true); + return null; + }); + } + + protected void handle(ChannelHandlerContext ctx, FullHttpRequest request) throws Exception { + String uri = request.uri(); + if (!uri.endsWith(mcpEndpoint)) { + sendError(ctx, HttpResponseStatus.NOT_FOUND, new McpError("Endpoint not found")); + return; + } + + if (isClosing.get()) { + sendError(ctx, HttpResponseStatus.SERVICE_UNAVAILABLE, new McpError("Server is shutting down")); + return; + } + + HttpMethod method = request.method(); + if (method == HttpMethod.GET) { + sendError(ctx, HttpResponseStatus.METHOD_NOT_ALLOWED, new McpError("GET method not allowed for stateless transport")); + } else if (method == HttpMethod.POST) { + handlePostRequest(ctx, request); + } else { + sendError(ctx, HttpResponseStatus.METHOD_NOT_ALLOWED, new McpError("Only POST method is supported")); + } + } + + /** + * Handles POST requests for incoming JSON-RPC messages from clients. + */ + private void handlePostRequest(ChannelHandlerContext ctx, FullHttpRequest request) { + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + + String accept = request.headers().get(ACCEPT); + if (accept == null || !(accept.contains(APPLICATION_JSON) && accept.contains(TEXT_EVENT_STREAM))) { + sendError(ctx, HttpResponseStatus.BAD_REQUEST, + new McpError("Both application/json and text/event-stream required in Accept header")); + return; + } + + try { + ByteBuf content = request.content(); + String body = content.toString(CharsetUtil.UTF_8); + + McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body); + + if (message instanceof McpSchema.JSONRPCRequest) { + McpSchema.JSONRPCRequest jsonrpcRequest = (McpSchema.JSONRPCRequest) message; + try { + this.mcpHandler.handleRequest(transportContext, jsonrpcRequest) + .thenAccept(jsonrpcResponse -> { + try { + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.OK, + Unpooled.copiedBuffer(objectMapper.writeValueAsString(jsonrpcResponse), CharsetUtil.UTF_8) + ); + + response.headers().set(HttpHeaderNames.CONTENT_TYPE, APPLICATION_JSON); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, response.content().readableBytes()); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + } catch (Exception e) { + logger.error("Failed to serialize response: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to serialize response: " + e.getMessage())); + } + }) + .exceptionally(e -> { + logger.error("Failed to handle request: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to handle request: " + e.getMessage())); + return null; + }); + } catch (Exception e) { + logger.error("Failed to handle request: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to handle request: " + e.getMessage())); + } + } else if (message instanceof McpSchema.JSONRPCNotification) { + McpSchema.JSONRPCNotification jsonrpcNotification = (McpSchema.JSONRPCNotification) message; + try { + this.mcpHandler.handleNotification(transportContext, jsonrpcNotification) + .thenRun(() -> { + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.ACCEPTED + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + }) + .exceptionally(e -> { + logger.error("Failed to handle notification: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to handle notification: " + e.getMessage())); + return null; + }); + } catch (Exception e) { + logger.error("Failed to handle notification: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to handle notification: " + e.getMessage())); + } + } else { + sendError(ctx, HttpResponseStatus.BAD_REQUEST, + new McpError("The server accepts either requests or notifications")); + } + } catch (IllegalArgumentException | IOException e) { + logger.error("Failed to deserialize message: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.BAD_REQUEST, new McpError("Invalid message format")); + } catch (Exception e) { + logger.error("Unexpected error handling message: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Unexpected error: " + e.getMessage())); + } + } + + /** + * Sends an error response to the client. + */ + private void sendError(ChannelHandlerContext ctx, HttpResponseStatus status, McpError mcpError) { + try { + String jsonError = objectMapper.writeValueAsString(mcpError); + ByteBuf content = Unpooled.copiedBuffer(jsonError, 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, response.content().readableBytes()); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + } catch (Exception e) { + logger.error(FAILED_TO_SEND_ERROR_RESPONSE, e.getMessage()); + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.INTERNAL_SERVER_ERROR + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + } + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStreamableHttpRequestHandler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStreamableHttpRequestHandler.java new file mode 100644 index 000000000..62025cbb9 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/handler/McpStreamableHttpRequestHandler.java @@ -0,0 +1,685 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +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.server.DefaultMcpTransportContext; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContext; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContextExtractor; +import com.taobao.arthas.mcp.server.protocol.spec.HttpHeaders; +import com.taobao.arthas.mcp.server.protocol.spec.*; +import com.taobao.arthas.mcp.server.util.Assert; +import com.taobao.arthas.mcp.server.util.KeepAliveScheduler; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelHandler; +import io.netty.channel.ChannelHandlerContext; +import io.netty.handler.codec.http.*; +import io.netty.util.CharsetUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Server-side implementation of the Model Context Protocol (MCP) streamable transport + * layer using HTTP with Server-Sent Events (SSE) through Netty. This implementation + * provides a bridge between Netty operations and the MCP transport interface. + * + * @see McpStreamableServerTransportProvider + */ +public class McpStreamableHttpRequestHandler { + + private static final Logger logger = LoggerFactory.getLogger(McpStreamableHttpRequestHandler.class); + + /** + * Event type for JSON-RPC messages sent through the SSE connection. + */ + public static final String MESSAGE_EVENT_TYPE = "message"; + + /** + * Header name for the response media types accepted by the requester. + */ + private static final String ACCEPT = "Accept"; + + public static final String UTF_8 = "UTF-8"; + public static final String APPLICATION_JSON = "application/json"; + public static final String TEXT_EVENT_STREAM = "text/event-stream"; + private static final String FAILED_TO_SEND_ERROR_RESPONSE = "Failed to send error response: {}"; + + /** + * The endpoint URI where clients should send their JSON-RPC messages. Defaults to + * "/mcp". + */ + private final String mcpEndpoint; + + /** + * Flag indicating whether DELETE requests are disallowed on the endpoint. + */ + private final boolean disallowDelete; + + private final ObjectMapper objectMapper; + + private McpStreamableServerSession.Factory sessionFactory; + + /** + * Map of active client sessions, keyed by mcp-session-id. + */ + private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + + private McpTransportContextExtractor contextExtractor; + + /** + * Flag indicating if the transport is shutting down. + */ + private final AtomicBoolean isClosing = new AtomicBoolean(false); + + /** + * Keep-alive scheduler for managing session pings. Activated if keepAliveInterval is + * set. Disabled by default. + */ + private KeepAliveScheduler keepAliveScheduler; + + /** + * Constructs a new NettyStreamableServerTransportProvider instance. + * + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of messages. + * @param mcpEndpoint The endpoint URI where clients should send their JSON-RPC + * messages via HTTP. + * @param disallowDelete Whether to disallow DELETE requests on the endpoint. + * @param contextExtractor The extractor for transport context from the request. + * @param keepAliveInterval Interval for keep-alive pings (null to disable) + * @throws IllegalArgumentException if any parameter is null + */ + public McpStreamableHttpRequestHandler(ObjectMapper objectMapper, String mcpEndpoint, + boolean disallowDelete, McpTransportContextExtractor contextExtractor, + Duration keepAliveInterval) { + this.objectMapper = objectMapper; + this.mcpEndpoint = mcpEndpoint; + this.disallowDelete = disallowDelete; + this.contextExtractor = contextExtractor; + + if (keepAliveInterval != null) { + this.keepAliveScheduler = KeepAliveScheduler + .builder(() -> this.isClosing.get() ? Collections.emptyList() : this.sessions.values()) + .initialDelay(keepAliveInterval) + .interval(keepAliveInterval) + .build(); + + this.keepAliveScheduler.start(); + logger.debug("Keep-alive scheduler started with interval: {}ms", keepAliveInterval.toMillis()); + } + } + + public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) { + this.sessionFactory = sessionFactory; + } + + public CompletableFuture notifyClients(String method, Object params) { + if (this.sessions.isEmpty()) { + logger.debug("No active sessions to broadcast message to"); + return CompletableFuture.completedFuture(null); + } + + logger.debug("Attempting to broadcast message to {} active sessions", this.sessions.size()); + + return CompletableFuture.runAsync(() -> { + this.sessions.values().parallelStream().forEach(session -> { + try { + session.sendNotification(method, params); + } catch (Exception e) { + logger.error("Failed to send message to session {}: {}", session.getId(), e.getMessage()); + } + }); + }); + } + + public CompletableFuture closeGracefully() { + return CompletableFuture.runAsync(() -> { + this.isClosing.set(true); + logger.debug("Initiating graceful shutdown with {} active sessions", this.sessions.size()); + + this.sessions.values().parallelStream().forEach(session -> { + try { + session.closeGracefully(); + } catch (Exception e) { + logger.error("Failed to close session {}: {}", session.getId(), e.getMessage()); + } + }); + + this.sessions.clear(); + logger.debug("Graceful shutdown completed"); + + if (this.keepAliveScheduler != null) { + this.keepAliveScheduler.shutdown(); + } + }); + } + + protected void handle(ChannelHandlerContext ctx, FullHttpRequest request) throws Exception { + String uri = request.uri(); + if (!uri.endsWith(mcpEndpoint)) { + sendError(ctx, HttpResponseStatus.NOT_FOUND, new McpError("Endpoint not found")); + return; + } + + if (isClosing.get()) { + sendError(ctx, HttpResponseStatus.SERVICE_UNAVAILABLE, new McpError("Server is shutting down")); + return; + } + + HttpMethod method = request.method(); + if (method == HttpMethod.GET) { + handleGetRequest(ctx, request); + } else if (method == HttpMethod.POST) { + handlePostRequest(ctx, request); + } else if (method == HttpMethod.DELETE) { + handleDeleteRequest(ctx, request); + } else { + sendError(ctx, HttpResponseStatus.METHOD_NOT_ALLOWED, new McpError("Method not allowed")); + } + } + + /** + * Handles GET requests to establish SSE connections and message replay. + */ + private void handleGetRequest(ChannelHandlerContext ctx, FullHttpRequest request) { + List badRequestErrors = new ArrayList<>(); + + String accept = request.headers().get(ACCEPT); + if (accept == null || !accept.contains(TEXT_EVENT_STREAM)) { + badRequestErrors.add("text/event-stream required in Accept header"); + } + + String sessionId = request.headers().get(HttpHeaders.MCP_SESSION_ID); + if (sessionId == null || sessionId.trim().isEmpty()) { + badRequestErrors.add("Session ID required in mcp-session-id header"); + } + + if (!badRequestErrors.isEmpty()) { + String combinedMessage = String.join("; ", badRequestErrors); + sendError(ctx, HttpResponseStatus.BAD_REQUEST, new McpError(combinedMessage)); + return; + } + + McpStreamableServerSession session = this.sessions.get(sessionId); + if (session == null) { + sendError(ctx, HttpResponseStatus.NOT_FOUND, new McpError("Session not found")); + return; + } + + logger.debug("Handling GET request for session: {}", sessionId); + + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + + try { + // Set up SSE response headers + 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, "keep-alive"); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + response.headers().set(HttpHeaderNames.TRANSFER_ENCODING, HttpHeaderValues.CHUNKED); + + ctx.writeAndFlush(response); + + NettyStreamableMcpSessionTransport sessionTransport = new NettyStreamableMcpSessionTransport( + sessionId, ctx); + + // Check if this is a replay request + String lastEventId = request.headers().get(HttpHeaders.LAST_EVENT_ID); + if (lastEventId != null) { + try { + // Replay messages from the last event ID + try { + session.replay(lastEventId).forEach(message -> { + try { + sessionTransport.sendMessage(message).join(); + } catch (Exception e) { + logger.error("Failed to replay message: {}", e.getMessage()); + ctx.close(); + } + }); + } catch (Exception e) { + logger.error("Failed to replay messages: {}", e.getMessage()); + ctx.close(); + } + } catch (Exception e) { + logger.error("Failed to replay messages: {}", e.getMessage()); + ctx.close(); + } + } else { + // Establish new listening stream + McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session + .listeningStream(sessionTransport); + + // Handle channel closure + ctx.channel().closeFuture().addListener(future -> { + logger.debug("SSE connection closed for session: {}", sessionId); + listeningStream.close(); + }); + } + } catch (Exception e) { + logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, new McpError("Internal server error")); + } + } + + /** + * Handles POST requests for incoming JSON-RPC messages from clients. + */ + private void handlePostRequest(ChannelHandlerContext ctx, FullHttpRequest request) { + List badRequestErrors = new ArrayList<>(); + + String accept = request.headers().get(ACCEPT); + if (accept == null || !accept.contains(TEXT_EVENT_STREAM)) { + badRequestErrors.add("text/event-stream required in Accept header"); + } + if (accept == null || !accept.contains(APPLICATION_JSON)) { + badRequestErrors.add("application/json required in Accept header"); + } + + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + + try { + ByteBuf content = request.content(); + String body = content.toString(CharsetUtil.UTF_8); + + McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body); + + // Handle initialization request + if (message instanceof McpSchema.JSONRPCRequest) { + McpSchema.JSONRPCRequest jsonrpcRequest = (McpSchema.JSONRPCRequest) message; + if (jsonrpcRequest.getMethod().equals(McpSchema.METHOD_INITIALIZE)) { + if (!badRequestErrors.isEmpty()) { + String combinedMessage = String.join("; ", badRequestErrors); + sendError(ctx, HttpResponseStatus.BAD_REQUEST, new McpError(combinedMessage)); + return; + } + + McpSchema.InitializeRequest initializeRequest = objectMapper.convertValue(jsonrpcRequest.getParams(), + new TypeReference() { + }); + McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory + .startSession(initializeRequest); + this.sessions.put(init.session().getId(), init.session()); + + try { + init.initResult() + .thenAccept(initResult -> { + try { + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.OK, + Unpooled.copiedBuffer(objectMapper.writeValueAsString( + new McpSchema.JSONRPCResponse( + McpSchema.JSONRPC_VERSION, jsonrpcRequest.getId(), initResult, null)), + CharsetUtil.UTF_8) + ); + + response.headers().set(HttpHeaderNames.CONTENT_TYPE, APPLICATION_JSON); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, response.content().readableBytes()); + response.headers().set(HttpHeaders.MCP_SESSION_ID, init.session().getId()); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + + ctx.writeAndFlush(response); + } catch (Exception e) { + logger.error("Failed to serialize init response: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to serialize response")); + } + }) + .exceptionally(e -> { + logger.error("Failed to initialize session: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to initialize session: " + e.getMessage())); + return null; + }); + return; + } catch (Exception e) { + logger.error("Failed to initialize session: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Failed to initialize session: " + e.getMessage())); + return; + } + } + } + + String sessionId = request.headers().get(HttpHeaders.MCP_SESSION_ID); + if (sessionId == null || sessionId.trim().isEmpty()) { + badRequestErrors.add("Session ID required in mcp-session-id header"); + } + + if (!badRequestErrors.isEmpty()) { + String combinedMessage = String.join("; ", badRequestErrors); + sendError(ctx, HttpResponseStatus.BAD_REQUEST, new McpError(combinedMessage)); + return; + } + + McpStreamableServerSession session = this.sessions.get(sessionId); + if (session == null) { + sendError(ctx, HttpResponseStatus.NOT_FOUND, + new McpError("Session not found: " + sessionId)); + return; + } + + if (message instanceof McpSchema.JSONRPCResponse) { + McpSchema.JSONRPCResponse jsonrpcResponse = (McpSchema.JSONRPCResponse) message; + session.accept(jsonrpcResponse) + .thenRun(() -> { + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.ACCEPTED + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + ctx.writeAndFlush(response); + }) + .exceptionally(e -> { + logger.error("Failed to accept response: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, new McpError("Failed to accept response")); + return null; + }); + } else if (message instanceof McpSchema.JSONRPCNotification) { + McpSchema.JSONRPCNotification jsonrpcNotification = (McpSchema.JSONRPCNotification) message; + session.accept(jsonrpcNotification, transportContext) + .thenRun(() -> { + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.ACCEPTED + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + ctx.writeAndFlush(response); + }) + .exceptionally(e -> { + logger.error("Failed to accept notification: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, new McpError("Failed to accept notification")); + return null; + }); + } else if (message instanceof McpSchema.JSONRPCRequest) { + McpSchema.JSONRPCRequest jsonrpcRequest = (McpSchema.JSONRPCRequest) message; + // For streaming responses, we need to return SSE + 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, "keep-alive"); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + response.headers().set(HttpHeaderNames.TRANSFER_ENCODING, HttpHeaderValues.CHUNKED); + + ctx.writeAndFlush(response); + + NettyStreamableMcpSessionTransport sessionTransport = new NettyStreamableMcpSessionTransport( + sessionId, ctx); + + try { + session.responseStream(jsonrpcRequest, sessionTransport, transportContext) + .exceptionally(e -> { + logger.error("Failed to handle request stream: {}", e.getMessage()); + ctx.close(); + return null; + }); + } catch (Exception e) { + logger.error("Failed to handle request stream: {}", e.getMessage()); + ctx.close(); + } + } else { + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Unknown message type")); + } + } catch (IllegalArgumentException | IOException e) { + logger.error("Failed to deserialize message: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.BAD_REQUEST, + new McpError("Invalid message format: " + e.getMessage())); + } catch (Exception e) { + logger.error("Error handling message: {}", e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Error processing message: " + e.getMessage())); + } + } + + /** + * Handles DELETE requests for session deletion. + */ + private void handleDeleteRequest(ChannelHandlerContext ctx, FullHttpRequest request) { + if (this.disallowDelete) { + sendError(ctx, HttpResponseStatus.METHOD_NOT_ALLOWED, new McpError("DELETE method not allowed")); + return; + } + + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + + String sessionId = request.headers().get(HttpHeaders.MCP_SESSION_ID); + if (sessionId == null) { + sendError(ctx, HttpResponseStatus.BAD_REQUEST, + new McpError("Session ID required in mcp-session-id header")); + return; + } + + McpStreamableServerSession session = this.sessions.get(sessionId); + if (session == null) { + sendError(ctx, HttpResponseStatus.NOT_FOUND, new McpError("Session not found")); + return; + } + + try { + session.delete() + .thenRun(() -> { + this.sessions.remove(sessionId); + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.OK + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + ctx.writeAndFlush(response); + }) + .exceptionally(e -> { + logger.error("Failed to delete session {}: {}", sessionId, e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError(e.getMessage())); + return null; + }); + } catch (Exception e) { + logger.error("Failed to delete session {}: {}", sessionId, e.getMessage()); + sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR, + new McpError("Error deleting session")); + } + } + + /** + * Sends an error response to the client. + */ + private void sendError(ChannelHandlerContext ctx, HttpResponseStatus status, McpError mcpError) { + try { + String jsonError = objectMapper.writeValueAsString(mcpError); + ByteBuf content = Unpooled.copiedBuffer(jsonError, 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, response.content().readableBytes()); + response.headers().set(HttpHeaderNames.ACCESS_CONTROL_ALLOW_ORIGIN, "*"); + + ctx.writeAndFlush(response); + } catch (Exception e) { + logger.error(FAILED_TO_SEND_ERROR_RESPONSE, e.getMessage()); + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.INTERNAL_SERVER_ERROR + ); + response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0); + ctx.writeAndFlush(response); + } + } + + // Implementation of McpStreamableServerTransport for Netty SSE sessions + private class NettyStreamableMcpSessionTransport implements McpStreamableServerTransport { + + private final String sessionId; + private final ChannelHandlerContext ctx; + private final AtomicBoolean closed = new AtomicBoolean(false); + private final ReentrantLock lock = new ReentrantLock(); + + NettyStreamableMcpSessionTransport(String sessionId, ChannelHandlerContext ctx) { + this.sessionId = sessionId; + this.ctx = ctx; + logger.debug("Streamable session transport {} initialized", sessionId); + } + + @Override + public CompletableFuture sendMessage(McpSchema.JSONRPCMessage message) { + return sendMessage(message, null); + } + + @Override + public CompletableFuture sendMessage(McpSchema.JSONRPCMessage message, String messageId) { + return CompletableFuture.runAsync(() -> { + if (this.closed.get()) { + logger.warn("Attempted to send message to closed session: {}", this.sessionId); + return; + } + + // Check if channel is still active + if (!this.ctx.channel().isActive()) { + logger.warn("Channel for session {} is not active, message will not be sent", this.sessionId); + return; + } + lock.lock(); + try { + if (this.closed.get()) { + logger.debug("Session {} was closed during message send attempt", this.sessionId); + return; + } + + String jsonText = objectMapper.writeValueAsString(message); + logger.debug("Sending SSE message to session {}: {}", this.sessionId, + jsonText.length() > 200 ? jsonText.substring(0, 200) + "..." : jsonText); + sendSseEvent(MESSAGE_EVENT_TYPE, jsonText, messageId != null ? messageId : this.sessionId); + logger.debug("Message sent to session {} with ID {}", this.sessionId, messageId); + } catch (Exception e) { + logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage()); + this.ctx.close(); + } finally { + lock.unlock(); + } + }); + } + + @Override + public T unmarshalFrom(Object data, TypeReference typeRef) { + return objectMapper.convertValue(data, typeRef); + } + + @Override + public CompletableFuture closeGracefully() { + return CompletableFuture.runAsync(this::close); + } + + @Override + public void close() { + lock.lock(); + try { + if (this.closed.get()) { + logger.debug("Session transport {} already closed", this.sessionId); + return; + } + + this.closed.set(true); + this.ctx.close(); + logger.debug("Successfully closed session transport {}", sessionId); + } catch (Exception e) { + logger.warn("Failed to close session transport {}: {}", sessionId, e.getMessage()); + } finally { + lock.unlock(); + } + } + + @Override + public Channel getChannel() { + return ctx.channel(); + } + + private void sendSseEvent(String eventType, String data, String id) { + StringBuilder sseData = new StringBuilder(); + if (id != null) { + sseData.append("id: ").append(id).append("\n"); + } + sseData.append("event: ").append(eventType).append("\n"); + sseData.append("data: ").append(data).append("\n\n"); + + ByteBuf buffer = Unpooled.copiedBuffer(sseData.toString(), CharsetUtil.UTF_8); + this.ctx.writeAndFlush(new DefaultHttpContent(buffer)); + + logger.debug("SSE event sent - Type: {}, ID: {}, Data length: {}", + eventType, id, data != null ? data.length() : 0); + } + } + + public static Builder builder() { + return new Builder(); + } + + /** + * Builder for creating instances of {@link McpStreamableHttpRequestHandler}. + */ + public static class Builder { + + private ObjectMapper objectMapper; + private String mcpEndpoint = "/mcp"; + private boolean disallowDelete = false; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + private Duration keepAliveInterval; + + public Builder objectMapper(ObjectMapper objectMapper) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + this.objectMapper = objectMapper; + return this; + } + + public Builder mcpEndpoint(String mcpEndpoint) { + Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); + this.mcpEndpoint = mcpEndpoint; + return this; + } + + public Builder disallowDelete(boolean disallowDelete) { + this.disallowDelete = disallowDelete; + return this; + } + + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "Context extractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + + public Builder keepAliveInterval(Duration keepAliveInterval) { + this.keepAliveInterval = keepAliveInterval; + return this; + } + + public McpStreamableHttpRequestHandler build() { + Assert.notNull(this.objectMapper, "ObjectMapper must be set"); + Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set"); + + return new McpStreamableHttpRequestHandler(this.objectMapper, this.mcpEndpoint, + this.disallowDelete, this.contextExtractor, this.keepAliveInterval); + } + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/store/InMemoryEventStore.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/store/InMemoryEventStore.java new file mode 100644 index 000000000..ab1bbadda --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/store/InMemoryEventStore.java @@ -0,0 +1,155 @@ +package com.taobao.arthas.mcp.server.protocol.server.store; + +import com.taobao.arthas.mcp.server.protocol.spec.EventStore; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +/** + * In-memory implementation of EventStore. + * + * @author Yeaury + */ +public class InMemoryEventStore implements EventStore { + + private static final Logger logger = LoggerFactory.getLogger(InMemoryEventStore.class); + + /** Global event ID counter */ + private final AtomicLong globalEventIdCounter = new AtomicLong(0); + + /** + * Events storage: sessionId -> list of events + */ + private final Map> sessionEvents = new ConcurrentHashMap<>(); + + /** + * Event ID to session mapping for fast lookup + */ + private final Map eventIdToSession = new ConcurrentHashMap<>(); + + /** + * Maximum events to keep per session (prevent memory leaks) + */ + private final int maxEventsPerSession; + + /** + * Default retention time in milliseconds (24 hours) + */ + private final long defaultRetentionMs; + + public InMemoryEventStore() { + this(1000, 24 * 60 * 60 * 1000L); // 1000 events, 24 hours + } + + public InMemoryEventStore(int maxEventsPerSession, long defaultRetentionMs) { + this.maxEventsPerSession = maxEventsPerSession; + this.defaultRetentionMs = defaultRetentionMs; + } + + @Override + public String storeEvent(String sessionId, McpSchema.JSONRPCMessage message) { + String eventId = String.valueOf(globalEventIdCounter.incrementAndGet()); + Instant timestamp = Instant.now(); + + StoredEvent event = new StoredEvent(eventId, sessionId, message, timestamp); + + sessionEvents.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(event); + eventIdToSession.put(eventId, sessionId); + + // Cleanup old events if needed + List events = sessionEvents.get(sessionId); + if (events.size() > maxEventsPerSession) { + // Remove oldest events + int toRemove = events.size() - maxEventsPerSession; + for (int i = 0; i < toRemove; i++) { + StoredEvent removedEvent = events.remove(0); + eventIdToSession.remove(removedEvent.getEventId()); + } + logger.debug("Cleaned up {} old events for session {}", toRemove, sessionId); + } + + logger.trace("Stored event {} for session {}", eventId, sessionId); + return eventId; + } + + @Override + public Stream getEventsForSession(String sessionId, String fromEventId) { + List events = sessionEvents.get(sessionId); + if (events == null || events.isEmpty()) { + return Stream.empty(); + } + + if (fromEventId == null) { + return Stream.empty(); + } + + boolean foundStartEvent = false; + List result = new ArrayList<>(); + + for (StoredEvent event : events) { + if (!foundStartEvent) { + if (event.getEventId().equals(fromEventId)) { + foundStartEvent = true; + result.add(event); + // clear the replayed events + events.remove(event); + eventIdToSession.remove(event.getEventId()); + } + continue; + } + result.add(event); + } + + return result.stream(); + } + + @Override + public void cleanupOldEvents(String sessionId, long maxAge) { + List events = sessionEvents.get(sessionId); + if (events == null || events.isEmpty()) { + return; + } + + Instant cutoff = Instant.now().minusMillis(maxAge); + + List toRemove = events.stream() + .filter(event -> event.getTimestamp().isBefore(cutoff)) + .collect(Collectors.toList()); + + for (StoredEvent event : toRemove) { + events.remove(event); + eventIdToSession.remove(event.getEventId()); + } + + if (!toRemove.isEmpty()) { + logger.debug("Cleaned up {} old events for session {}", toRemove.size(), sessionId); + } + } + + @Override + public void removeSessionEvents(String sessionId) { + List events = sessionEvents.remove(sessionId); + if (events != null) { + for (StoredEvent event : events) { + eventIdToSession.remove(event.getEventId()); + } + logger.debug("Removed {} events for session {}", events.size(), sessionId); + } + } + + public void cleanupExpiredEvents() { + for (String sessionId : sessionEvents.keySet()) { + cleanupOldEvents(sessionId, defaultRetentionMs); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/HttpNettyServerTransportProvider.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/HttpNettyServerTransportProvider.java deleted file mode 100644 index a21259dca..000000000 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/HttpNettyServerTransportProvider.java +++ /dev/null @@ -1,90 +0,0 @@ -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 notifyClients(String method, Object params) { - return mcpRequestHandler.notifyClients(method, params); - } - - @Override - public CompletableFuture 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(); - } - -} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStatelessServerTransport.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStatelessServerTransport.java new file mode 100644 index 000000000..42e0186a9 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStatelessServerTransport.java @@ -0,0 +1,112 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server.transport; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.taobao.arthas.mcp.server.protocol.server.McpStatelessServerHandler; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContextExtractor; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpStatelessHttpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.spec.McpStatelessServerTransport; +import com.taobao.arthas.mcp.server.util.Assert; +import io.netty.handler.codec.http.FullHttpRequest; + +import java.util.concurrent.CompletableFuture; + +/** + * Server-side implementation of the Model Context Protocol (MCP) stateless transport + * layer using HTTP through Netty. This implementation provides a bridge between + * Netty operations and the MCP transport interface for stateless operations. + * + * @see McpStatelessServerTransport + */ +public class NettyStatelessServerTransport implements McpStatelessServerTransport { + + private final McpStatelessHttpRequestHandler requestHandler; + + /** + * Constructs a new NettyStatelessServerTransportProvider instance. + * + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of messages. + * @param mcpEndpoint The endpoint URI where clients should send their JSON-RPC + * messages via HTTP. + * @param contextExtractor The extractor for transport context from the request. + * @throws IllegalArgumentException if any parameter is null + */ + private NettyStatelessServerTransport(ObjectMapper objectMapper, String mcpEndpoint, + McpTransportContextExtractor contextExtractor) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); + + this.requestHandler = new McpStatelessHttpRequestHandler(objectMapper, mcpEndpoint, contextExtractor); + } + + @Override + public void setMcpHandler(McpStatelessServerHandler mcpHandler) { + requestHandler.setMcpHandler(mcpHandler); + } + + /** + * Initiates a graceful shutdown of the transport. + * + * @return A CompletableFuture that completes when all cleanup operations are finished + */ + @Override + public CompletableFuture closeGracefully() { + return requestHandler.closeGracefully(); + } + + /** + * Gets the underlying HTTP request handler. + * + * @return The McpStatelessHttpRequestHandler instance + */ + public McpStatelessHttpRequestHandler getMcpRequestHandler() { + if (this.requestHandler != null) { + return this.requestHandler; + } + throw new UnsupportedOperationException("Stateless transport provider does not support request handler"); + } + + public static Builder builder() { + return new Builder(); + } + + /** + * Builder for creating instances of {@link NettyStatelessServerTransport}. + */ + public static class Builder { + + private ObjectMapper objectMapper; + private String mcpEndpoint = "/mcp"; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + + public Builder objectMapper(ObjectMapper objectMapper) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + this.objectMapper = objectMapper; + return this; + } + + public Builder mcpEndpoint(String mcpEndpoint) { + Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); + this.mcpEndpoint = mcpEndpoint; + return this; + } + + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "Context extractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + + public NettyStatelessServerTransport build() { + Assert.notNull(this.objectMapper, "ObjectMapper must be set"); + Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set"); + + return new NettyStatelessServerTransport(this.objectMapper, this.mcpEndpoint, this.contextExtractor); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStreamableServerTransportProvider.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStreamableServerTransportProvider.java new file mode 100644 index 000000000..28a06df62 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/server/transport/NettyStreamableServerTransportProvider.java @@ -0,0 +1,148 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.server.transport; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContextExtractor; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpStreamableHttpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.spec.McpStreamableServerSession; +import com.taobao.arthas.mcp.server.protocol.spec.McpStreamableServerTransportProvider; +import com.taobao.arthas.mcp.server.util.Assert; +import io.netty.handler.codec.http.FullHttpRequest; + +import java.time.Duration; +import java.util.concurrent.CompletableFuture; + +/** + * Server-side implementation of the Model Context Protocol (MCP) streamable transport + * layer using HTTP with Server-Sent Events (SSE) through Netty. This implementation + * provides a bridge between Netty operations and the MCP transport interface. + * + * @see McpStreamableServerTransportProvider + */ +public class NettyStreamableServerTransportProvider implements McpStreamableServerTransportProvider { + + private final McpStreamableHttpRequestHandler requestHandler; + + /** + * Constructs a new NettyStreamableServerTransportProvider instance. + * + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of messages. + * @param mcpEndpoint The endpoint URI where clients should send their JSON-RPC + * messages via HTTP. + * @param disallowDelete Whether to disallow DELETE requests on the endpoint. + * @param contextExtractor The extractor for transport context from the request. + * @param keepAliveInterval Interval for keep-alive pings (null to disable) + * @throws IllegalArgumentException if any parameter is null + */ + private NettyStreamableServerTransportProvider(ObjectMapper objectMapper, String mcpEndpoint, + boolean disallowDelete, McpTransportContextExtractor contextExtractor, + Duration keepAliveInterval) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); + + this.requestHandler = new McpStreamableHttpRequestHandler(objectMapper, mcpEndpoint, disallowDelete, contextExtractor, keepAliveInterval); + } + + @Override + public String protocolVersion() { + return "2025-03-26"; + } + + @Override + public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) { + requestHandler.setSessionFactory(sessionFactory); + } + + /** + * Broadcasts a notification to all connected clients through their SSE connections. + * + * @param method The method name for the notification + * @param params The parameters for the notification + * @return A CompletableFuture that completes when the broadcast attempt is finished + */ + @Override + public CompletableFuture notifyClients(String method, Object params) { + return requestHandler.notifyClients(method, params); + } + + /** + * Initiates a graceful shutdown of the transport. + * + * @return A CompletableFuture that completes when all cleanup operations are finished + */ + @Override + public CompletableFuture closeGracefully() { + return requestHandler.closeGracefully(); + } + + @Override + public McpStreamableHttpRequestHandler getMcpRequestHandler() { + if (this.requestHandler != null) { + return this.requestHandler; + } + throw new UnsupportedOperationException("Streamable transport provider does not support legacy SSE request handler"); + } + + public static Builder builder() { + return new Builder(); + } + + /** + * Builder for creating instances of {@link NettyStreamableServerTransportProvider}. + */ + public static class Builder { + + private ObjectMapper objectMapper; + private String mcpEndpoint = "/mcp"; + private boolean disallowDelete = false; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + private Duration keepAliveInterval; + + public Builder objectMapper(ObjectMapper objectMapper) { + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + this.objectMapper = objectMapper; + return this; + } + + public Builder mcpEndpoint(String mcpEndpoint) { + Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); + this.mcpEndpoint = mcpEndpoint; + return this; + } + + public Builder disallowDelete(boolean disallowDelete) { + this.disallowDelete = disallowDelete; + return this; + } + + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "Context extractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + + public Builder keepAliveInterval(Duration keepAliveInterval) { + this.keepAliveInterval = keepAliveInterval; + return this; + } + + public NettyStreamableServerTransportProvider build() { + Assert.notNull(this.objectMapper, "ObjectMapper must be set"); + Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set"); + + return new NettyStreamableServerTransportProvider(this.objectMapper, this.mcpEndpoint, + this.disallowDelete, this.contextExtractor, this.keepAliveInterval); + } + } + + // Placeholder interface for KeepAliveScheduler if not already implemented + private interface KeepAliveScheduler { + void shutdown(); + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/DefaultMcpStreamableServerSessionFactory.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/DefaultMcpStreamableServerSessionFactory.java new file mode 100644 index 000000000..8becafec9 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/DefaultMcpStreamableServerSessionFactory.java @@ -0,0 +1,64 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +import com.taobao.arthas.mcp.server.CommandExecutor; +import com.taobao.arthas.mcp.server.protocol.server.McpInitRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.McpNotificationHandler; +import com.taobao.arthas.mcp.server.protocol.server.McpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.store.InMemoryEventStore; + +import java.time.Duration; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; + +/** + * Default implementation of the streamable server session factory. + * This factory creates new MCP streamable server sessions with the provided configuration. + * + */ +public class DefaultMcpStreamableServerSessionFactory implements McpStreamableServerSession.Factory { + + private final Duration requestTimeout; + private final McpInitRequestHandler mcpInitRequestHandler; + private final Map> requestHandlers; + private final Map notificationHandlers; + private final CommandExecutor commandExecutor; + + public DefaultMcpStreamableServerSessionFactory(Duration requestTimeout, + McpInitRequestHandler mcpInitRequestHandler, + Map> requestHandlers, + Map notificationHandlers, + CommandExecutor commandExecutor) { + this.requestTimeout = requestTimeout; + this.mcpInitRequestHandler = mcpInitRequestHandler; + this.requestHandlers = requestHandlers; + this.notificationHandlers = notificationHandlers; + this.commandExecutor = commandExecutor; + } + + @Override + public McpStreamableServerSession.McpStreamableServerSessionInit startSession( + McpSchema.InitializeRequest initializeRequest) { + + // Create a new session with a unique ID + McpStreamableServerSession session = new McpStreamableServerSession( + UUID.randomUUID().toString(), + initializeRequest.getCapabilities(), + initializeRequest.getClientInfo(), + requestTimeout, + requestHandlers, + notificationHandlers, + commandExecutor, + new InMemoryEventStore()); + + // Handle the initialization request + CompletableFuture initResult = + this.mcpInitRequestHandler.handle(initializeRequest); + + return new McpStreamableServerSession.McpStreamableServerSessionInit(session, initResult); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/EventStore.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/EventStore.java new file mode 100644 index 000000000..39e23d5d1 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/EventStore.java @@ -0,0 +1,51 @@ +package com.taobao.arthas.mcp.server.protocol.spec; + +import java.time.Instant; +import java.util.List; +import java.util.function.Consumer; +import java.util.stream.Stream; + +/** + * EventStore interface for storing and replaying JSON-RPC events. + * Supports event persistence and stream resumability for MCP Streamable HTTP protocol. + * + * @author Yeaury + */ +public interface EventStore { + + class StoredEvent { + private final String eventId; + private final String sessionId; + private final McpSchema.JSONRPCMessage message; + private final Instant timestamp; + + public StoredEvent(String eventId, String sessionId, McpSchema.JSONRPCMessage message, Instant timestamp) { + this.eventId = eventId; + this.sessionId = sessionId; + this.message = message; + this.timestamp = timestamp; + } + + public String getEventId() { + return eventId; + } + public String getSessionId() { + return sessionId; + } + public McpSchema.JSONRPCMessage getMessage() { + return message; + } + public Instant getTimestamp() { + return timestamp; + } + } + + String storeEvent(String sessionId, McpSchema.JSONRPCMessage message); + + Stream getEventsForSession(String sessionId, String fromEventId); + + void cleanupOldEvents(String sessionId, long maxAge); + + void removeSessionEvents(String sessionId); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/HttpHeaders.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/HttpHeaders.java new file mode 100644 index 000000000..4e13b076a --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/HttpHeaders.java @@ -0,0 +1,24 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +public interface HttpHeaders { + + /** + * Identifies individual MCP sessions. + */ + String MCP_SESSION_ID = "mcp-session-id"; + + /** + * Identifies events within an SSE Stream. + */ + String LAST_EVENT_ID = "last-event-id"; + + /** + * Identifies the MCP protocol version. + */ + String PROTOCOL_VERSION = "MCP-Protocol-Version"; + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpError.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpError.java index 2c93d646e..b128477bb 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpError.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpError.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.spec; import com.taobao.arthas.mcp.server.protocol.spec.McpSchema.JSONRPCResponse.JSONRPCError; diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSchema.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSchema.java index ab7b5ca64..f1122f66d 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSchema.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSchema.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.spec; import com.fasterxml.jackson.annotation.*; @@ -9,11 +13,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.util.*; /** * Based on the JSON-RPC 2.0 specification @@ -29,7 +29,7 @@ public final class McpSchema { private McpSchema() { } - public static final String LATEST_PROTOCOL_VERSION = "2024-11-05"; + public static final String LATEST_PROTOCOL_VERSION = "2025-03-26"; public static final String JSONRPC_VERSION = "2.0"; @@ -44,6 +44,8 @@ public final class McpSchema { public static final String METHOD_PING = "ping"; + public static final String METHOD_NOTIFICATION_PROGRESS = "notifications/progress"; + // Tool Methods public static final String METHOD_TOOLS_LIST = "tools/list"; @@ -82,15 +84,6 @@ public final class McpSchema { private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - // Streamable Methods - public static final String METHOD_PROGRESS_UPDATE = "notifications/progress"; - - public static final String METHOD_STATUS_UPDATE = "notifications/status"; - - public static final String METHOD_INTERMEDIATE_RESULT = "notifications/intermediate"; - - public static final String METHOD_TOOL_ERROR = "notifications/error"; - // --------------------------- // JSON-RPC Error Codes // --------------------------- @@ -1302,16 +1295,24 @@ public final class McpSchema { public static class CallToolRequest implements Request { private final String name; private final Map arguments; + private final Map meta; + public CallToolRequest( @JsonProperty("name") String name, - @JsonProperty("arguments") Map arguments) { + @JsonProperty("arguments") Map arguments, + @JsonProperty("_meta") Map meta) { this.name = name; this.arguments = arguments; + this.meta = meta; } public CallToolRequest(String name, String jsonArguments) { - this(name, parseJsonArguments(jsonArguments)); + this(name, parseJsonArguments(jsonArguments), null); + } + + public CallToolRequest(String name, Map arguments) { + this(name, arguments, null); } private static Map parseJsonArguments(String jsonArguments) { @@ -1330,6 +1331,56 @@ public final class McpSchema { public Map getArguments() { return arguments; } + + public Map getMeta() { + return meta; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + + private String name; + + private Map arguments; + + private Map meta; + + public Builder name(String name) { + this.name = name; + return this; + } + + public Builder arguments(Map arguments) { + this.arguments = arguments; + return this; + } + + public Builder arguments(String jsonArguments) { + this.arguments = parseJsonArguments(jsonArguments); + return this; + } + + public Builder meta(Map meta) { + this.meta = meta; + return this; + } + + public Builder progressToken(String progressToken) { + if (this.meta == null) { + this.meta = new HashMap<>(); + } + this.meta.put("progressToken", progressToken); + return this; + } + + public CallToolRequest build() { + Assert.hasText(name, "name must not be empty"); + return new CallToolRequest(name, arguments, meta); + } + } } /** @@ -1661,12 +1712,12 @@ public final class McpSchema { public static class LoggingMessageNotification { private final LoggingLevel level; private final String logger; - private final String data; + private final Object data; public LoggingMessageNotification( @JsonProperty("level") LoggingLevel level, @JsonProperty("logger") String logger, - @JsonProperty("data") String data) { + @JsonProperty("data") Object data) { this.level = level; this.logger = logger; this.data = data; @@ -1680,7 +1731,7 @@ public final class McpSchema { return logger; } - public String getData() { + public Object getData() { return data; } @@ -1691,7 +1742,7 @@ public final class McpSchema { public static class Builder { private LoggingLevel level = LoggingLevel.INFO; private String logger = "server"; - private String data; + private Object data; public Builder level(LoggingLevel level) { this.level = level; @@ -1703,7 +1754,7 @@ public final class McpSchema { return this; } - public Builder data(String data) { + public Builder data(Object data) { this.data = data; return this; } @@ -1750,6 +1801,32 @@ public final class McpSchema { } } + /** + * Notification for sending intermediate results during streaming tool execution. + * This allows tools to send partial results to clients in real-time. + */ + @JsonIgnoreProperties(ignoreUnknown = true) + public static class IntermediateResultNotification { + private final String type; + private final Object data; + + public IntermediateResultNotification( + @JsonProperty("type") String type, + @JsonProperty("data") Object data) { + this.type = type; + this.data = data; + } + + public String getType() { + return type; + } + + public Object getData() { + return data; + } + + } + // --------------------------- // Autocomplete // --------------------------- diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerSession.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerSession.java deleted file mode 100644 index 234fced1c..000000000 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerSession.java +++ /dev/null @@ -1,433 +0,0 @@ -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> 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> requestHandlers; - - private final Map notificationHandlers; - - private final McpServerTransport transport; - - private final CompletableFuture exchangeFuture = new CompletableFuture<>(); - - private final AtomicReference clientCapabilities = new AtomicReference<>(); - - private final AtomicReference 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> requestHandlers, Map 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 CompletableFuture sendRequest(String method, Object requestParams, TypeReference typeRef) { - String requestId = this.generateRequestId(); - - CompletableFuture result = new CompletableFuture<>(); - CompletableFuture 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 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 handle(McpSchema.JSONRPCMessage message) { - CompletableFuture 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 result) { - logger.debug("Received response: {}", response); - CompletableFuture 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 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 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 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() { - }); - 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 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 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 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 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 typedHandler = (RequestHandler) 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 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 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 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 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 data = new HashMap<>(); - data.put("method", method); - return new MethodNotFoundError(method, "method not found: " + method, data); - } - - @Override - public CompletableFuture closeGracefully() { - return this.transport.closeGracefully(); - } - - @Override - public void close() { - this.transport.close(); - } - - public interface InitRequestHandler { - - CompletableFuture handle(McpSchema.InitializeRequest initializeRequest); - - } - - public interface InitNotificationHandler { - - CompletableFuture handle(); - - } - - public interface NotificationHandler { - - CompletableFuture handle(McpNettyServerExchange exchange, Object params); - - } - - public interface RequestHandler { - - CompletableFuture handle(McpNettyServerExchange exchange, Object params); - } - - @FunctionalInterface - public interface Factory { - - McpServerSession create(McpServerTransport sessionTransport); - } - -} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransport.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransport.java index 742b3eaf1..74b63443a 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransport.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.spec; import io.netty.channel.Channel; diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransportProvider.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransportProvider.java index 342d9d840..4e2f3a012 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransportProvider.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpServerTransportProvider.java @@ -1,6 +1,10 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.spec; -import com.taobao.arthas.mcp.server.protocol.server.handler.McpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.handler.McpStreamableHttpRequestHandler; import java.util.concurrent.CompletableFuture; @@ -12,8 +16,6 @@ import java.util.concurrent.CompletableFuture; */ public interface McpServerTransportProvider { - void setSessionFactory(McpServerSession.Factory sessionFactory); - CompletableFuture notifyClients(String method, Object params); CompletableFuture closeGracefully(); @@ -22,5 +24,7 @@ public interface McpServerTransportProvider { closeGracefully(); } - McpRequestHandler getMcpRequestHandler(); + McpStreamableHttpRequestHandler getMcpRequestHandler(); + + String protocolVersion(); } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSession.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSession.java index 14b682f42..686855dd6 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSession.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.spec; import java.util.concurrent.CompletableFuture; diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStatelessServerTransport.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStatelessServerTransport.java new file mode 100644 index 000000000..57f0b1b51 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStatelessServerTransport.java @@ -0,0 +1,21 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +import com.taobao.arthas.mcp.server.protocol.server.McpStatelessServerHandler; + +import java.util.concurrent.CompletableFuture; + +public interface McpStatelessServerTransport { + + void setMcpHandler(McpStatelessServerHandler mcpHandler); + + default void close() { + this.closeGracefully(); + } + + CompletableFuture closeGracefully(); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerSession.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerSession.java new file mode 100644 index 000000000..619fcaf3c --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerSession.java @@ -0,0 +1,428 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.taobao.arthas.mcp.server.CommandExecutor; +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; + +import com.taobao.arthas.mcp.server.protocol.server.McpNotificationHandler; +import com.taobao.arthas.mcp.server.protocol.server.McpRequestHandler; +import com.taobao.arthas.mcp.server.protocol.server.McpTransportContext; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.session.ArthasCommandSessionManager; +import com.taobao.arthas.mcp.server.util.Assert; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; +import java.util.stream.Stream; + +/** + * Main implementation of a streamable MCP server session that manages client connections + * and handles JSON-RPC communication using CompletableFuture for async operations. + */ +public class McpStreamableServerSession implements McpSession { + + private static final Logger logger = LoggerFactory.getLogger(McpStreamableServerSession.class); + + private final ConcurrentHashMap requestIdToStream = new ConcurrentHashMap<>(); + + private final String id; + private final Duration requestTimeout; + private final AtomicLong requestCounter = new AtomicLong(0); + private final Map> requestHandlers; + private final Map notificationHandlers; + + private final AtomicReference clientCapabilities = new AtomicReference<>(); + + private final AtomicReference clientInfo = new AtomicReference<>(); + + private final AtomicReference listeningStreamRef; + + private final MissingMcpTransportSession missingMcpTransportSession; + + private volatile McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.INFO; + + private final CommandExecutor commandExecutor; + + private final ArthasCommandSessionManager commandSessionManager; + + private final EventStore eventStore; + + public McpStreamableServerSession(String id, McpSchema.ClientCapabilities clientCapabilities, + McpSchema.Implementation clientInfo, Duration requestTimeout, + Map> requestHandlers, + Map notificationHandlers, + CommandExecutor commandExecutor, EventStore eventStore) { + this.id = id; + this.missingMcpTransportSession = new MissingMcpTransportSession(id); + this.listeningStreamRef = new AtomicReference<>(this.missingMcpTransportSession); + this.clientCapabilities.lazySet(clientCapabilities); + this.clientInfo.lazySet(clientInfo); + this.requestTimeout = requestTimeout; + this.requestHandlers = requestHandlers; + this.notificationHandlers = notificationHandlers; + this.commandExecutor = commandExecutor; + this.commandSessionManager = new ArthasCommandSessionManager(commandExecutor); + this.eventStore = eventStore; + } + + /** + * Sets the minimum logging level for this session. + * @param minLoggingLevel the minimum logging level + */ + public void setMinLoggingLevel(McpSchema.LoggingLevel minLoggingLevel) { + Assert.notNull(minLoggingLevel, "minLoggingLevel must not be null"); + this.minLoggingLevel = minLoggingLevel; + } + + /** + * Checks if notifications for the given logging level are allowed. + * @param loggingLevel the logging level to check + * @return true if notifications for this level are allowed + */ + public boolean isNotificationForLevelAllowed(McpSchema.LoggingLevel loggingLevel) { + return loggingLevel.level() >= this.minLoggingLevel.level(); + } + + public String getId() { + return this.id; + } + + private String generateRequestId() { + return this.id + "-" + this.requestCounter.getAndIncrement(); + } + + @Override + public CompletableFuture sendRequest(String method, Object requestParams, TypeReference typeRef) { + McpSession listeningStream = this.listeningStreamRef.get(); + return listeningStream.sendRequest(method, requestParams, typeRef); + } + + @Override + public CompletableFuture sendNotification(String method, Object params) { + McpSession listeningStream = this.listeningStreamRef.get(); + return listeningStream.sendNotification(method, params); + } + + public CompletableFuture delete() { + return this.closeGracefully().thenRun(() -> { + try { + eventStore.removeSessionEvents(this.id); + commandSessionManager.closeCommandSession(this.id); + } catch (Exception e) { + logger.warn("Failed to clear session during deletion: {}", e.getMessage()); + } + }); + } + + public McpStreamableServerSessionStream listeningStream(McpStreamableServerTransport transport) { + McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport); + this.listeningStreamRef.set(listeningStream); + return listeningStream; + } + + /** + * 重播会话事件，从指定的最后事件ID之后开始 + * + * @param lastEventId 最后一个事件ID，如果为null则从头开始重播 + * @return 事件消息流 + */ + public Stream replay(Object lastEventId) { + String lastEventIdStr = lastEventId != null ? lastEventId.toString() : null; + + return eventStore.getEventsForSession(this.id, lastEventIdStr) + .map(EventStore.StoredEvent::getMessage); + } + + public CompletableFuture responseStream(McpSchema.JSONRPCRequest jsonrpcRequest, + McpStreamableServerTransport transport, McpTransportContext transportContext) { + + McpStreamableServerSessionStream stream = new McpStreamableServerSessionStream(transport); + McpRequestHandler requestHandler = this.requestHandlers.get(jsonrpcRequest.getMethod()); + + if (requestHandler == null) { + MethodNotFoundError error = getMethodNotFoundError(jsonrpcRequest.getMethod()); + McpSchema.JSONRPCResponse errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + jsonrpcRequest.getId(), null, + new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND, + error.getMessage(), error.getData())); + + // 存储错误响应 + try { + eventStore.storeEvent(this.id, errorResponse); + } catch (Exception e) { + logger.warn("Failed to store error response event: {}", e.getMessage()); + } + + return transport.sendMessage(errorResponse, null); + } + ArthasCommandContext commandContext = createCommandContext(); + + return requestHandler + .handle(new McpNettyServerExchange(this.id, stream, clientCapabilities.get(), + clientInfo.get(), transportContext), commandContext, jsonrpcRequest.getParams()) + .thenApply(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + jsonrpcRequest.getId(), result, null)) + .thenCompose(response -> transport.sendMessage(response, null)) + .thenCompose(v -> transport.closeGracefully()); + } + + public CompletableFuture accept(McpSchema.JSONRPCNotification notification, + McpTransportContext transportContext) { + + McpNotificationHandler notificationHandler = this.notificationHandlers.get(notification.getMethod()); + if (notificationHandler == null) { + logger.error("No handler registered for notification method: {}", notification.getMethod()); + return CompletableFuture.completedFuture(null); + } + + ArthasCommandContext commandContext = createCommandContext(); + McpSession listeningStream = this.listeningStreamRef.get(); + return notificationHandler.handle(new McpNettyServerExchange(this.id, listeningStream, + this.clientCapabilities.get(), this.clientInfo.get(), transportContext), commandContext, notification.getParams()); + } + + public CompletableFuture accept(McpSchema.JSONRPCResponse response) { + McpStreamableServerSessionStream stream = this.requestIdToStream.get(response.getId()); + if (stream == null) { + CompletableFuture f = CompletableFuture.completedFuture(null); + f.completeExceptionally(new McpError("Unexpected response for unknown id " + response.getId())); + return f; + } + + CompletableFuture future = stream.pendingResponses.remove(response.getId()); + if (future == null) { + CompletableFuture f = CompletableFuture.completedFuture(null); + f.completeExceptionally(new McpError("Unexpected response for unknown id " + response.getId())); + return f; + } else { + future.complete(response); + } + + return CompletableFuture.completedFuture(null); + } + + public 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; + } + } + + + private MethodNotFoundError getMethodNotFoundError(String method) { + return new MethodNotFoundError(method, "Method not found: " + method, null); + } + + @Override + public CompletableFuture closeGracefully() { + McpSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession); + + // 清理 Arthas 命令会话 + try { + commandSessionManager.closeCommandSession(this.id); + logger.debug("Successfully closed command session during graceful shutdown: {}", this.id); + } catch (Exception e) { + logger.warn("Failed to close command session during graceful shutdown: {}", e.getMessage()); + } + + return listeningStream.closeGracefully(); + // TODO: Also close all the open streams + } + + @Override + public void close() { + McpSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession); + + // 清理 Arthas 命令会话 + try { + commandSessionManager.closeCommandSession(this.id); + logger.debug("Successfully closed command session during close: {}", this.id); + } catch (Exception e) { + logger.warn("Failed to close command session during close: {}", e.getMessage()); + } + + if (listeningStream != null) { + listeningStream.close(); + } + // TODO: Also close all open streams + } + + public interface Factory { + McpStreamableServerSessionInit startSession(McpSchema.InitializeRequest initializeRequest); + } + + public static class McpStreamableServerSessionInit { + private final McpStreamableServerSession session; + private final CompletableFuture initResult; + + public McpStreamableServerSessionInit( + McpStreamableServerSession session, + CompletableFuture initResult) { + this.session = session; + this.initResult = initResult; + } + + public McpStreamableServerSession session() { + return session; + } + + public CompletableFuture initResult() { + return initResult; + } + } + + + public final class McpStreamableServerSessionStream implements McpSession { + + private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); + + private final McpStreamableServerTransport transport; + private final String transportId; + private final Supplier uuidGenerator; + + public McpStreamableServerSessionStream(McpStreamableServerTransport transport) { + this.transport = transport; + this.transportId = UUID.randomUUID().toString(); + // This ID design allows for a constant-time extraction of the history by + // precisely identifying the SSE stream using the first component + this.uuidGenerator = () -> this.transportId + "_" + UUID.randomUUID(); + } + + @Override + public CompletableFuture sendRequest(String method, Object requestParams, TypeReference typeRef) { + String requestId = McpStreamableServerSession.this.generateRequestId(); + + McpStreamableServerSession.this.requestIdToStream.put(requestId, this); + + CompletableFuture responseFuture = new CompletableFuture<>(); + this.pendingResponses.put(requestId, responseFuture); + + McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, + method, requestId, requestParams); + String messageId = null; + + // 存储发送的请求到事件存储 + try { + messageId = McpStreamableServerSession.this.eventStore.storeEvent( + McpStreamableServerSession.this.id, jsonrpcRequest); + } catch (Exception e) { + logger.warn("Failed to store outbound request event: {}", e.getMessage()); + } + + // Send the message + this.transport.sendMessage(jsonrpcRequest, messageId).exceptionally(ex -> { + responseFuture.completeExceptionally(ex); + return null; + }); + + return responseFuture.handle((jsonRpcResponse, throwable) -> { + // Cleanup + this.pendingResponses.remove(requestId); + McpStreamableServerSession.this.requestIdToStream.remove(requestId); + + if (throwable != null) { + if (throwable instanceof RuntimeException) { + throw (RuntimeException) throwable; + } + throw new RuntimeException(throwable); + } + + if (jsonRpcResponse.getError() != null) { + throw new RuntimeException(new McpError(jsonRpcResponse.getError())); + } else { + if (typeRef.getType().equals(Void.class)) { + return null; + } else { + return this.transport.unmarshalFrom(jsonRpcResponse.getResult(), typeRef); + } + } + }); + } + + @Override + public CompletableFuture sendNotification(String method, Object params) { + McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification( + McpSchema.JSONRPC_VERSION, method, params); + String messageId = null; + try { + messageId = McpStreamableServerSession.this.eventStore.storeEvent( + McpStreamableServerSession.this.id, jsonrpcNotification); + } catch (Exception e) { + logger.warn("Failed to store outbound notification event: {}", e.getMessage()); + } + + return this.transport.sendMessage(jsonrpcNotification, messageId); + } + + @Override + public CompletableFuture closeGracefully() { + // Complete all pending responses with error + this.pendingResponses.values().forEach(future -> + future.completeExceptionally(new RuntimeException("Stream closed"))); + this.pendingResponses.clear(); + + // If this was the generic stream, reset it + McpStreamableServerSession.this.listeningStreamRef.compareAndSet(this, + McpStreamableServerSession.this.missingMcpTransportSession); + + McpStreamableServerSession.this.requestIdToStream.values().removeIf(this::equals); + + return this.transport.closeGracefully(); + } + + @Override + public void close() { + this.pendingResponses.values().forEach(future -> + future.completeExceptionally(new RuntimeException("Stream closed"))); + this.pendingResponses.clear(); + + // If this was the generic stream, reset it + McpStreamableServerSession.this.listeningStreamRef.compareAndSet(this, + McpStreamableServerSession.this.missingMcpTransportSession); + McpStreamableServerSession.this.requestIdToStream.values().removeIf(this::equals); + + this.transport.close(); + } + + + } + + /** + * 创建命令执行上下文 + * + * @return 命令执行上下文 + */ + private ArthasCommandContext createCommandContext() { + ArthasCommandSessionManager.CommandSessionBinding binding = commandSessionManager.getCommandSession(this.id); + return new ArthasCommandContext(commandExecutor, binding); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransport.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransport.java new file mode 100644 index 000000000..c6699d1ee --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransport.java @@ -0,0 +1,21 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +import com.fasterxml.jackson.core.type.TypeReference; + +import java.util.concurrent.CompletableFuture; + +/** + * Marker interface for the server-side MCP streamable transport. + * This extends the basic server transport with streamable message sending capabilities. + * + */ +public interface McpStreamableServerTransport extends McpServerTransport { + + CompletableFuture sendMessage(McpSchema.JSONRPCMessage message, String messageId); + + T unmarshalFrom(Object value, TypeReference typeRef); +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransportProvider.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransportProvider.java new file mode 100644 index 000000000..7e518530d --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpStreamableServerTransportProvider.java @@ -0,0 +1,22 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +import java.util.concurrent.CompletableFuture; + +public interface McpStreamableServerTransportProvider extends McpServerTransportProvider { + + + void setSessionFactory(McpStreamableServerSession.Factory sessionFactory); + + CompletableFuture notifyClients(String method, Object params); + + default void close() { + this.closeGracefully(); + } + + String protocolVersion(); + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpTransport.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpTransport.java index ffe8b1820..cfa93c3fc 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpTransport.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/McpTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package com.taobao.arthas.mcp.server.protocol.spec; import java.util.concurrent.CompletableFuture; diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/MissingMcpTransportSession.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/MissingMcpTransportSession.java new file mode 100644 index 000000000..079d4e58a --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/protocol/spec/MissingMcpTransportSession.java @@ -0,0 +1,52 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.protocol.spec; + +import com.fasterxml.jackson.core.type.TypeReference; + +import java.util.concurrent.CompletableFuture; + +/** + * A placeholder implementation of McpLoggableSession that represents a missing or unavailable transport session. + * This class is used when no active transport is available but session operations are attempted. + */ +public class MissingMcpTransportSession implements McpSession { + + private final String sessionId; + + public MissingMcpTransportSession(String sessionId) { + this.sessionId = sessionId; + } + + @Override + public CompletableFuture sendRequest(String method, Object requestParams, TypeReference typeRef) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally( + new IllegalStateException("Stream unavailable for session " + this.sessionId) + ); + return future; + } + + @Override + public CompletableFuture sendNotification(String method, Object params) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally( + new IllegalStateException("Stream unavailable for session " + this.sessionId) + ); + return future; + } + + @Override + public CompletableFuture closeGracefully() { + return CompletableFuture.completedFuture(null); + } + + @Override + public void close() { + // Nothing to close for a missing session + } + + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandContext.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandContext.java new file mode 100644 index 000000000..19239d375 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandContext.java @@ -0,0 +1,76 @@ +package com.taobao.arthas.mcp.server.session; + +import com.taobao.arthas.mcp.server.CommandExecutor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.Objects; + +/** + * 命令执行上下文 + * + * @author Yeaury + */ +public class ArthasCommandContext { + + private final CommandExecutor commandExecutor; + + private final ArthasCommandSessionManager.CommandSessionBinding sessionBinding; + + private static final long DEFAULT_SYNC_TIMEOUT = 10000L; + + public ArthasCommandContext(CommandExecutor commandExecutor) { + this.commandExecutor = Objects.requireNonNull(commandExecutor, "commandExecutor cannot be null"); + this.sessionBinding = null; + } + + public ArthasCommandContext(CommandExecutor commandExecutor, ArthasCommandSessionManager.CommandSessionBinding sessionBinding) { + this.commandExecutor = Objects.requireNonNull(commandExecutor, "commandExecutor cannot be null"); + this.sessionBinding = Objects.requireNonNull(sessionBinding, "sessionBinding cannot be null"); + } + + private void requireSessionSupport() { + if (sessionBinding == null) { + throw new IllegalStateException("Session-based operations are not supported in temporary mode. " + + "Use ArthasCommandContext(CommandExecutor, CommandSessionBinding) constructor to enable session support."); + } + } + + public boolean isSessionSupported() { + return sessionBinding != null; + } + + public Map executeSync(String commandLine) { + return executeSync(commandLine, DEFAULT_SYNC_TIMEOUT); + } + + public Map executeSync(String commandLine, long timeout) { + return commandExecutor.executeSync(commandLine, timeout); + } + + public Map executeSyncWithSession(String commandLine, long timeout) { + requireSessionSupport(); + return commandExecutor.executeSync(commandLine, timeout, sessionBinding.getArthasSessionId()); + } + + public Map executeAsync(String commandLine) { + requireSessionSupport(); + return commandExecutor.executeAsync(commandLine, sessionBinding.getArthasSessionId()); + } + + public Map pullResults() { + requireSessionSupport(); + return commandExecutor.pullResults(sessionBinding.getArthasSessionId(), sessionBinding.getConsumerId()); + } + + public Map interruptJob() { + requireSessionSupport(); + return commandExecutor.interruptJob(sessionBinding.getArthasSessionId()); + } + + public String getArthasSessionId() { + requireSessionSupport(); + return sessionBinding.getArthasSessionId(); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java new file mode 100644 index 000000000..abe78f0eb --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java @@ -0,0 +1,137 @@ +package com.taobao.arthas.mcp.server.session; + +import com.taobao.arthas.mcp.server.CommandExecutor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Arthas命令执行Session管理器 + * + * @author Yeaury + */ +public class ArthasCommandSessionManager { + + private static final Logger logger = LoggerFactory.getLogger(ArthasCommandSessionManager.class); + + private final CommandExecutor commandExecutor; + private final ConcurrentHashMap sessionBindings = new ConcurrentHashMap<>(); + + public ArthasCommandSessionManager(CommandExecutor commandExecutor) { + this.commandExecutor = commandExecutor; + } + + public static class CommandSessionBinding { + private final String mcpSessionId; + private final String arthasSessionId; + private final String consumerId; + + public CommandSessionBinding(String mcpSessionId, String arthasSessionId, String consumerId) { + this.mcpSessionId = mcpSessionId; + this.arthasSessionId = arthasSessionId; + this.consumerId = consumerId; + } + + public String getMcpSessionId() { + return mcpSessionId; + } + + public String getArthasSessionId() { + return arthasSessionId; + } + + public String getConsumerId() { + return consumerId; + } + } + + public CommandSessionBinding createCommandSession(String mcpSessionId) { + Map result = commandExecutor.createSession(); + + CommandSessionBinding binding = new CommandSessionBinding( + mcpSessionId, + (String) result.get("sessionId"), + (String) result.get("consumerId") + ); + + return binding; + } + + public CommandSessionBinding getCommandSession(String mcpSessionId) { + CommandSessionBinding binding = sessionBindings.get(mcpSessionId); + + if (binding == null) { + binding = createCommandSession(mcpSessionId); + sessionBindings.put(mcpSessionId, binding); + logger.debug("Created new command session: MCP={}, Arthas={}", mcpSessionId, binding.getArthasSessionId()); + return binding; + } + + if (isSessionValid(binding)) { + logger.debug("Using existing valid session: MCP={}, Arthas={}", mcpSessionId, binding.getArthasSessionId()); + return binding; + } + + logger.info("Session expired, recreating: MCP={}, Arthas={}", mcpSessionId, binding.getArthasSessionId()); + + try { + commandExecutor.closeSession(binding.getArthasSessionId()); + } catch (Exception e) { + logger.debug("Failed to close expired session (may already be cleaned up): {}", e.getMessage()); + } + + CommandSessionBinding newBinding = createCommandSession(mcpSessionId); + sessionBindings.put(mcpSessionId, newBinding); + logger.info("Recreated command session: MCP={}, Old Arthas={}, New Arthas={}", + mcpSessionId, binding.getArthasSessionId(), newBinding.getArthasSessionId()); + + return newBinding; + } + + /** + * 检查session是否有效 + * 通过尝试获取结果来验证session和consumer是否仍然存在 + */ + private boolean isSessionValid(CommandSessionBinding binding) { + try { + Map result = commandExecutor.pullResults(binding.getArthasSessionId(), binding.getConsumerId()); + + if (result == null) { + return true; + } + + Boolean success = (Boolean) result.get("success"); + if (Boolean.TRUE.equals(success)) { + return true; + } + + String errorMessage = (String) result.get("error"); + if (errorMessage != null && (errorMessage.contains("Session not found") || + errorMessage.contains("Consumer not found") || + errorMessage.contains("session is inactive"))) { + logger.debug("Session validation failed: {}", errorMessage); + return false; + } + + return true; + + } catch (Exception e) { + logger.debug("Session validation error: {}", e.getMessage()); + return false; + } + } + + public void closeCommandSession(String mcpSessionId) { + CommandSessionBinding binding = sessionBindings.remove(mcpSessionId); + if (binding != null) { + commandExecutor.closeSession(binding.getArthasSessionId()); + logger.debug("Closed command session: MCP={}, Arthas={}", mcpSessionId, binding.getArthasSessionId()); + } + } + + public void closeAllSessions() { + sessionBindings.keySet().forEach(this::closeCommandSession); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/DefaultToolCallback.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/DefaultToolCallback.java index 33f3bb7fb..d115ab218 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/DefaultToolCallback.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/DefaultToolCallback.java @@ -13,10 +13,12 @@ import org.slf4j.LoggerFactory; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.lang.reflect.Modifier; +import java.lang.reflect.Parameter; import java.lang.reflect.Type; import java.util.Arrays; import java.util.Map; import java.util.stream.Stream; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; public class DefaultToolCallback implements ToolCallback { @@ -62,6 +64,8 @@ public class DefaultToolCallback implements ToolCallback { validateToolContextSupport(toolContext); Map toolArguments = extractToolArguments(toolInput); + + validateRequiredParameters(toolArguments); Object[] methodArguments = buildMethodArguments(toolArguments, toolContext); @@ -85,6 +89,34 @@ public class DefaultToolCallback implements ToolCallback { } } + /** + * validate the required parameters + */ + private void validateRequiredParameters(Map toolArguments) { + Parameter[] parameters = this.toolMethod.getParameters(); + + for (Parameter parameter : parameters) { + if (parameter.getType().isAssignableFrom(ToolContext.class)) { + continue; + } + + ToolParam toolParam = parameter.getAnnotation(ToolParam.class); + if (toolParam != null && toolParam.required()) { + String paramName = parameter.getName(); + Object paramValue = toolArguments.get(paramName); + + // check if the parameter is empty or an empty string + if (paramValue == null) { + throw new IllegalArgumentException("Required parameter '" + paramName + "' is missing"); + } + + if (paramValue instanceof String && ((String) paramValue).trim().isEmpty()) { + throw new IllegalArgumentException("Required parameter '" + paramName + "' cannot be empty"); + } + } + } + } + private Map extractToolArguments(String toolInput) { return JsonParser.fromJson(toolInput, new TypeReference>() { }); diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/definition/ToolDefinition.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/definition/ToolDefinition.java index 2788b99b3..1ebaad310 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/definition/ToolDefinition.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/definition/ToolDefinition.java @@ -31,6 +31,10 @@ public class ToolDefinition { return inputSchema; } + public boolean isStreamable() { + return streamable; + } + public static Builder builder() { return new Builder(); } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/ArthasCommandExecutor.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/ArthasCommandExecutor.java deleted file mode 100644 index 2cb3bf930..000000000 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/ArthasCommandExecutor.java +++ /dev/null @@ -1,54 +0,0 @@ -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 result = executor.execute(command, timeout); - return JsonParser.toJson(result); - } catch (Exception e) { - logger.error("Failed to execute command: {}", command, e); - Map errorResponse = new HashMap<>(); - errorResponse.put("error", true); - errorResponse.put("message", "Error executing command '" + command + "': " + e.getMessage()); - return JsonParser.toJson(errorResponse); - } - } - -} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/StreamableToolUtils.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/StreamableToolUtils.java new file mode 100644 index 000000000..21eeeec9d --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/StreamableToolUtils.java @@ -0,0 +1,262 @@ +package com.taobao.arthas.mcp.server.tool.function; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * 同步工具的工具类 + * 提供同步执行命令并收集所有结果的功能 + * + * @author Yeaury + */ +public final class StreamableToolUtils { + + private static final Logger logger = LoggerFactory.getLogger(StreamableToolUtils.class); + + private static final int DEFAULT_POLL_INTERVAL_MS = 200; // 默认轮询间隔200ms + private static final int ERROR_RETRY_INTERVAL_MS = 500; // 错误重试间隔500ms + + private static final int MAX_POLL_ATTEMPTS = 100; // 最大轮询尝试次数，避免无限循环 + private static final int MAX_ERROR_RETRIES = 10; // 最大错误重试次数 + + public static final int MIN_ALLOW_INPUT_COUNT_TO_COMPLETE = 2; + + private StreamableToolUtils() { + } + + /** + * 同步执行命令并收集所有结果，支持进度通知 + * + * @param exchange MCP交换器，用于发送进度通知 + * @param commandContext 命令上下文 + * @param expectedResultCount 预期结果数量 + * @param intervalMs 轮询间隔 + * @param progressToken 进度令牌 + * @return 包含所有结果的Map，如果执行失败返回null + */ + public static Map executeAndCollectResults(McpNettyServerExchange exchange, + ArthasCommandContext commandContext, + Integer expectedResultCount, Integer intervalMs, + String progressToken) { + List allResults = new ArrayList<>(); + int pollAttempts = 0; + int errorRetries = 0; + int allowInputCount = 0; + int totalResultCount = 0; + + // 轮询间隔使用命令执行间隙的 1/10,事件驱动则在命令中自定义默认轮询间隔 + // 工具中默认轮询间隔为200ms + int pullIntervalMs = (intervalMs != null && intervalMs > 0) ? intervalMs : DEFAULT_POLL_INTERVAL_MS; + + try { + while (pollAttempts < MAX_POLL_ATTEMPTS) { + pollAttempts++; + + try { + Map results = commandContext.pullResults(); + if (results == null) { + Thread.sleep(pullIntervalMs); + continue; + } + errorRetries = 0; + + Map filteredResults = filterCommandSpecificResults(results); + List currentBatchResults = getCommandSpecificResults(filteredResults); + + if (currentBatchResults != null && !currentBatchResults.isEmpty()) { + allResults.addAll(currentBatchResults); + totalResultCount += currentBatchResults.size(); + logger.debug("Collected {} results, total: {}", currentBatchResults.size(), totalResultCount); + + if (exchange != null) { + sendProgressNotification(exchange, totalResultCount, + expectedResultCount != null ? expectedResultCount : totalResultCount, + progressToken); + } + } + + boolean commandCompleted = checkCommandCompletion(results, allowInputCount); + if (commandCompleted) { + allowInputCount++; + } + + String jobStatus = (String) results.get("jobStatus"); + + // 判断是否应该结束 + // 如果是TERMINATED状态，或者命令已完成且允许输入次数大于等于2，或者实际结果数量达到预期结果数量 + if ("TERMINATED".equals(jobStatus) + || (commandCompleted && allowInputCount >= MIN_ALLOW_INPUT_COUNT_TO_COMPLETE) + || (expectedResultCount != null && totalResultCount >= expectedResultCount)) { + logger.info("Command completed. Total results collected: {}, Expected: {}", totalResultCount, expectedResultCount); + break; + } + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + logger.warn("Command execution interrupted"); + return null; + } catch (Exception e) { + if (++errorRetries >= MAX_ERROR_RETRIES) { + logger.error("Maximum error retries exceeded", e); + return null; + } + + try { + Thread.sleep(ERROR_RETRY_INTERVAL_MS); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + return null; + } + } + } + + Map finalResult = new HashMap<>(); + finalResult.put("results", allResults); + finalResult.put("resultCount", totalResultCount); + finalResult.put("status", "completed"); + finalResult.put("stage", "final"); + + if (pollAttempts >= MAX_POLL_ATTEMPTS) { + logger.warn("Command execution reached maximum poll attempts: {}", MAX_POLL_ATTEMPTS); + finalResult.put("warning", "Execution reached maximum poll attempts"); + } + + return finalResult; + + } catch (Exception e) { + logger.error("Error in command execution", e); + return null; + } + } + + private static boolean checkCommandCompletion(Map results, int currentAllowInputCount) { + if (results == null) { + return false; + } + + @SuppressWarnings("unchecked") + List resultList = (List) results.get("results"); + if (resultList == null || resultList.isEmpty()) { + return false; + } + + for (Object result : resultList) { + if (result.getClass().getSimpleName().equals("InputStatusModel")) { + try { + java.lang.reflect.Method getInputStatusMethod = result.getClass().getMethod("getInputStatus"); + Object inputStatusObj = getInputStatusMethod.invoke(result); + if (inputStatusObj != null) { + String inputStatusName = inputStatusObj.getClass().getSimpleName().equals("InputStatus") ? + inputStatusObj.toString() : String.valueOf(inputStatusObj); + if ("ALLOW_INPUT".equals(inputStatusName)) { + logger.debug("Command completion detected via InputStatusModel: ALLOW_INPUT (count: {})", currentAllowInputCount + 1); + return true; + } + } + } catch (Exception e) { + logger.debug("Failed to get inputStatus from InputStatusModel", e); + } + } + } + + return false; + } + + private static Map filterCommandSpecificResults(Map results) { + if (results == null) { + return new HashMap<>(); + } + + Map filteredResults = new HashMap<>(results); + @SuppressWarnings("unchecked") + List resultList = (List) results.get("results"); + + if (resultList == null || resultList.isEmpty()) { + return filteredResults; + } + + // 定义需要排除的辅助模型类型 + String[] auxiliaryModelTypes = { + "InputStatusModel", "StatusModel", "WelcomeModel", "MessageModel", + "CommandRequestModel", "SessionModel", "EnhancerModel" + }; + + List filteredResultList = resultList.stream() + .filter(result -> { + String resultClassName = result.getClass().getSimpleName(); + + for (String auxiliaryType : auxiliaryModelTypes) { + if (resultClassName.equals(auxiliaryType)) { + return false; + } + } + + return true; + }) + .collect(java.util.stream.Collectors.toList()); + + filteredResults.put("results", filteredResultList); + filteredResults.put("resultCount", filteredResultList.size()); + + return filteredResults; + } + + private static List getCommandSpecificResults(Map filteredResults) { + if (filteredResults == null) { + return new ArrayList<>(); + } + + @SuppressWarnings("unchecked") + List resultList = (List) filteredResults.get("results"); + return resultList != null ? resultList : new ArrayList<>(); + } + + /** + * 发送进度通知 + */ + private static void sendProgressNotification(McpNettyServerExchange exchange, int currentResultCount, + int totalExpected, String progressToken) { + try { + if (progressToken != null) { + exchange.progressNotification(new McpSchema.ProgressNotification( + progressToken, + currentResultCount, + (double) totalExpected + )).join(); + } + + } catch (Exception e) { + logger.error("Error sending progress notification", e); + } + } + + public static Map createErrorResponse(String message) { + Map response = new HashMap<>(); + response.put("error", true); + response.put("message", message); + response.put("status", "error"); + response.put("stage", "final"); + return response; + } + + public static Map createCompletedResponse(String message, Map results) { + Map response = new HashMap<>(); + response.put("status", "completed"); + response.put("message", message); + response.put("stage", "final"); + + if (results != null) { + response.putAll(results); + } + + return response; + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/DashboardTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/DashboardTool.java new file mode 100644 index 000000000..018d95e10 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/DashboardTool.java @@ -0,0 +1,81 @@ +package com.taobao.arthas.mcp.server.tool.function.jvm300; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class DashboardTool { + + private static final Logger logger = LoggerFactory.getLogger(DashboardTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_REFRESH_INTERVAL_MS = 3000; + + /** + * dashboard 实时面板命令 + * 支持: + * - intervalMs: 刷新间隔，单位 ms，默认 3000ms + * - count: 刷新次数限制，即 -n 参数；如果不指定则使用 DEFAULT_NUMBER_OF_EXECUTIONS (3次) + */ + @Tool( + name = "dashboard", + description = "Dashboard 诊断工具: 实时展示 JVM/应用面板，可利用参数控制诊断次数与间隔。对应 Arthas 的 dashboard 命令。", + streamable = true + ) + public String dashboard( + @ToolParam(description = "刷新间隔，单位为毫秒，默认 3000ms。用于控制输出频率", required = false) + Integer intervalMs, + + @ToolParam(description = "执行次数限制，默认值为 3。达到指定次数后自动停止", required = false) + Integer numberOfExecutions, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + int interval = (intervalMs != null && intervalMs > 0) ? intervalMs : DEFAULT_REFRESH_INTERVAL_MS; + int execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + + try { + StringBuilder cmd = new StringBuilder("dashboard"); + cmd.append(" -i ").append(interval); + cmd.append(" -n ").append(execCount); + String commandStr = cmd.toString(); + logger.info("Starting dashboard execution: {}", commandStr); + + // 如果只执行一次，使用同步执行 + if (execCount == 1) { + logger.info("Executing sync dashboard command: {}", commandStr); + Map result = commandContext.executeSync(commandStr); + return JsonParser.toJson(result); + } else { + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, interval / 10, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("Dashboard execution completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("Dashboard execution failed due to timeout or error limits exceeded")); + } + } + + } catch (Exception e) { + logger.error("Error executing dashboard command", e); + return JsonParser.toJson(createErrorResponse("Error executing dashboard: " + e.getMessage())); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/GetStaticTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/GetStaticTool.java index 9432df9a0..61f43e306 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/GetStaticTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/GetStaticTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class GetStaticTool { @@ -25,8 +29,11 @@ public class GetStaticTool { String fieldName, @ToolParam(description = "OGNL 表达式", required = false) - String ognlExpression + String ognlExpression, + + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("getstatic"); if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) { @@ -43,6 +50,6 @@ public class GetStaticTool { } String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/HeapdumpTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/HeapdumpTool.java new file mode 100644 index 000000000..1bb8a788d --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/HeapdumpTool.java @@ -0,0 +1,70 @@ +package com.taobao.arthas.mcp.server.tool.function.jvm300; + +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; + +import java.time.LocalDateTime; +import java.time.format.DateTimeFormatter; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; + +public class HeapdumpTool { + + public static final String DEFAULT_DUMP_DIR = Paths.get("arthas-output").toAbsolutePath().toString().replace("\\", "/"); + + private static final DateTimeFormatter TIMESTAMP_FORMATTER = DateTimeFormatter.ofPattern("yyyyMMdd_HHmmss"); + + /** + * heapdump 诊断工具 + * 支持: + * - live: 是否只 dump 存活对象 (--live) + * - filePath: 输出文件路径，若为空则使用默认临时文件 + */ + @Tool( + name = "heapdump", + description = "Heapdump 诊断工具: 生成 JVM heap dump，支持 --live 选项。对应 Arthas 的 heapdump 命令。" + ) + public String heapdump( + @ToolParam(description = "是否只 dump 存活对象 (--live)", required = false) + Boolean live, + + @ToolParam(description = "指定输出文件路径，默认为当前工作目录下的arthas-output文件夹中的时间戳命名的.hprof文件", required = false) + String filePath, + + ToolContext toolContext + ) throws IOException { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + + String finalFilePath; + + if (filePath != null && !filePath.trim().isEmpty()) { + finalFilePath = filePath.trim().replace("\\", "/"); + } else { + Path defaultDir = Paths.get(DEFAULT_DUMP_DIR); + if (!Files.exists(defaultDir)) { + Files.createDirectories(defaultDir); + } + + String timestamp = LocalDateTime.now().format(TIMESTAMP_FORMATTER); + String defaultFileName = String.format("heapdump_%s.hprof", timestamp); + finalFilePath = Paths.get(DEFAULT_DUMP_DIR, defaultFileName).toString().replace("\\", "/"); + } + + StringBuilder cmd = new StringBuilder("heapdump"); + if (Boolean.TRUE.equals(live)) { + cmd.append(" --live"); + } + cmd.append(" ").append(finalFilePath); + + String commandStr = cmd.toString(); + return JsonParser.toJson(commandContext.executeSync(commandStr)); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/JvmTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/JvmTool.java index 4e3516cd8..2bee437ee 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/JvmTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/JvmTool.java @@ -1,7 +1,11 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; import com.taobao.arthas.mcp.server.tool.annotation.Tool; -import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class JvmTool { @@ -9,8 +13,9 @@ public class JvmTool { name = "jvm", description = "Jvm 诊断工具: 查看当前 JVM 运行时信息。对应 Arthas 的 jvm 命令。" ) - public String jvm() { + public String jvm(ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); String commandStr = "jvm"; - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MBeanTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MBeanTool.java new file mode 100644 index 000000000..de785b8f2 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MBeanTool.java @@ -0,0 +1,123 @@ +package com.taobao.arthas.mcp.server.tool.function.jvm300; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class MBeanTool { + + private static final Logger logger = LoggerFactory.getLogger(MBeanTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_REFRESH_INTERVAL_MS = 3000; + + /** + * mbean 诊断工具: 查看或监控 MBean 属性 + * 支持: + * - namePattern: MBean 名称表达式，支持通配符或正则（需开启 -E） + * - attributePattern: 属性名称表达式，支持通配符或正则（需开启 -E） + * - metadata: 是否查看元信息 (-m) + * - intervalMs: 刷新属性值时间间隔 (ms) (-i)，required=false + * - numberOfExecutions: 刷新次数 (-n)，若未指定或 <=0 则使用 DEFAULT_NUMBER_OF_EXECUTIONS + * - regex: 是否启用正则匹配 (-E)，required=false + */ + @Tool( + name = "mbean", + description = "MBean 诊断工具: 查看或监控 MBean 属性信息，对应 Arthas 的 mbean 命令。", + streamable = true + ) + public String mbean( + @ToolParam(description = "MBean名称表达式匹配，如java.lang:type=GarbageCollector,name=*") + String namePattern, + + @ToolParam(description = "属性名表达式匹配，支持通配符如CollectionCount", required = false) + String attributePattern, + + @ToolParam(description = "是否查看元信息 (-m)", required = false) + Boolean metadata, + + @ToolParam(description = "刷新间隔，单位为毫秒，默认 3000ms。用于控制输出频率", required = false) + Integer intervalMs, + + @ToolParam(description = "执行次数限制，默认值为 3。达到指定次数后自动停止", required = false) + Integer numberOfExecutions, + + @ToolParam(description = "开启正则表达式匹配，默认为通配符匹配，默认false", required = false) + Boolean regex, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + boolean needStreamOutput = (intervalMs != null && intervalMs > 0) || (numberOfExecutions != null && numberOfExecutions > 0); + + int interval = DEFAULT_REFRESH_INTERVAL_MS; + int execCount = DEFAULT_NUMBER_OF_EXECUTIONS; + + if (needStreamOutput) { + interval = (intervalMs != null && intervalMs > 0) ? intervalMs : DEFAULT_REFRESH_INTERVAL_MS; + execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + } + + try { + StringBuilder cmd = new StringBuilder("mbean"); + + if (Boolean.TRUE.equals(metadata)) { + cmd.append(" -m"); + } + if (Boolean.TRUE.equals(regex)) { + cmd.append(" -E"); + } + + // 只有在需要流式输出且不是查看元数据时才添加 -i 和 -n 参数 + if (needStreamOutput && !Boolean.TRUE.equals(metadata)) { + cmd.append(" -i ").append(interval); + cmd.append(" -n ").append(execCount); + } + + if (namePattern != null && !namePattern.trim().isEmpty()) { + cmd.append(" ").append(namePattern.trim()); + } + if (attributePattern != null && !attributePattern.trim().isEmpty()) { + cmd.append(" ").append(attributePattern.trim()); + } + + String commandStr = cmd.toString(); + logger.info("Starting mbean execution: {}", commandStr); + + // 使用同步执行的情况：查看元数据 或者 不需要流式输出 + if (Boolean.TRUE.equals(metadata) || !needStreamOutput) { + logger.info("Executing sync mbean command: {}", commandStr); + Map result = commandContext.executeSync(commandStr); + return JsonParser.toJson(result); + } else { + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, interval / 10, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("MBean execution completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("MBean execution failed due to timeout or error limits exceeded")); + } + } + + } catch (Exception e) { + logger.error("Error executing mbean command", e); + return JsonParser.toJson(createErrorResponse("Error executing mbean: " + e.getMessage())); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MemoryTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MemoryTool.java index 755c4b365..1fe178d6a 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MemoryTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/MemoryTool.java @@ -1,7 +1,11 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; import com.taobao.arthas.mcp.server.tool.annotation.Tool; -import com.taobao.arthas.mcp.server.tool.function.ArthasCommandExecutor; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class MemoryTool { @@ -9,8 +13,9 @@ public class MemoryTool { name = "memory", description = "Memory 诊断工具: 查看 JVM 内存使用情况，对应 Arthas 的 memory 命令。" ) - public String memory() { + public String memory(ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); String commandStr = "memory"; - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/OgnlTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/OgnlTool.java index 551c53951..8b33fc6c5 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/OgnlTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/OgnlTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class OgnlTool { @@ -21,8 +25,11 @@ public class OgnlTool { String classLoaderClass, @ToolParam(description = "结果对象展开层次 (-x)，默认 1", required = false) - Integer expandLevel + Integer expandLevel, + + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("ognl"); if (classLoaderHash != null && !classLoaderHash.trim().isEmpty()) { @@ -37,6 +44,6 @@ public class OgnlTool { cmd.append(" ").append(expression.trim()); String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/PerfCounterTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/PerfCounterTool.java index 2f61c0b49..4b8490ef7 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/PerfCounterTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/PerfCounterTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class PerfCounterTool { @@ -12,13 +16,15 @@ public class PerfCounterTool { ) public String perfcounter( @ToolParam(description = "是否打印更多详情 (-d)", required = false) - Boolean detailed + Boolean detailed, + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("perfcounter"); if (Boolean.TRUE.equals(detailed)) { cmd.append(" -d"); } String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysEnvTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysEnvTool.java index aa2fac2c5..8138530bd 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysEnvTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysEnvTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class SysEnvTool { @@ -12,13 +16,15 @@ public class SysEnvTool { ) public String sysenv( @ToolParam(description = "环境变量名。若为空或空字符串，则查看所有变量；否则查看单个变量值。", required = false) - String envName + String envName, + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("sysenv"); if (envName != null && !envName.trim().isEmpty()) { cmd.append(" ").append(envName.trim()); } String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysPropTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysPropTool.java index 52e170552..ef03ccb09 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysPropTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/SysPropTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class SysPropTool { @@ -15,8 +19,11 @@ public class SysPropTool { String propertyName, @ToolParam(description = "属性值；若指定则修改，否则查看", required = false) - String propertyValue + String propertyValue, + + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("sysprop"); if (propertyName != null && !propertyName.trim().isEmpty()) { cmd.append(" ").append(propertyName.trim()); @@ -25,6 +32,6 @@ public class SysPropTool { } } String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/ThreadTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/ThreadTool.java new file mode 100644 index 000000000..1715f4f76 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/ThreadTool.java @@ -0,0 +1,68 @@ +package com.taobao.arthas.mcp.server.tool.function.jvm300; + +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; + +public class ThreadTool { + + private static final Logger logger = LoggerFactory.getLogger(ThreadTool.class); + + /** + * thread 诊断工具: 查看线程信息及堆栈 + * 支持: + * - threadId: 线程 ID，required=false + * - topN: 最忙前 N 个线程并打印堆栈 (-n)，required=false + * - blocking: 是否查找阻塞其他线程的线程 (-b)，required=false + * - all: 是否显示所有匹配线程 (--all)，required=false + */ + @Tool( + name = "thread", + description = "Thread 诊断工具: 查看线程信息及堆栈，对应 Arthas 的 thread 命令。一次性输出结果。" + ) + public String thread( + @ToolParam(description = "线程 ID", required = false) + Long threadId, + + @ToolParam(description = "最忙前 N 个线程并打印堆栈 (-n)", required = false) + Integer topN, + + @ToolParam(description = "是否查找阻塞其他线程的线程 (-b)", required = false) + Boolean blocking, + + @ToolParam(description = "是否显示所有匹配线程 (--all)", required = false) + Boolean all, + + ToolContext toolContext + ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + + StringBuilder cmd = new StringBuilder("thread"); + + if (Boolean.TRUE.equals(blocking)) { + cmd.append(" -b"); + } + if (topN != null && topN > 0) { + cmd.append(" -n ").append(topN); + } + if (Boolean.TRUE.equals(all)) { + cmd.append(" --all"); + } + if (threadId != null && threadId > 0) { + cmd.append(" ").append(threadId); + } + + String commandStr = cmd.toString(); + logger.info("Executing thread command: {}", commandStr); + + return JsonParser.toJson(commandContext.executeSync(commandStr)); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMOptionTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMOptionTool.java index fce7cd99c..5097dd1b5 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMOptionTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMOptionTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class VMOptionTool { @@ -15,8 +19,11 @@ public class VMOptionTool { String key, @ToolParam(description = "更新值，仅在更新时使用", required = false) - String value + String value, + + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("vmoption"); if (key != null && !key.trim().isEmpty()) { cmd.append(" ").append(key.trim()); @@ -25,6 +32,6 @@ public class VMOptionTool { } } String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMToolTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMToolTool.java index 6b6c6a985..96cbbaab9 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMToolTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/jvm300/VMToolTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.jvm300; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class VMToolTool { @@ -37,8 +41,11 @@ public class VMToolTool { String express, @ToolParam(description = "线程 ID (-t)，interruptThread 时使用", required = false) - Long threadId + Long threadId, + + ToolContext toolContext ) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("vmtool"); if (action == null || action.trim().isEmpty()) { throw new IllegalArgumentException("vmtool: action 参数不能为空"); @@ -76,7 +83,7 @@ public class VMToolTool { } String commandStr = cmd.toString(); - return ArthasCommandExecutor.executeCommand(commandStr); + return JsonParser.toJson(commandContext.executeSync(commandStr)); } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/ClassLoaderTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/ClassLoaderTool.java index 9c2180938..5722326b2 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/ClassLoaderTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/ClassLoaderTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class ClassLoaderTool { @@ -30,7 +34,10 @@ public class ClassLoaderTool { String resource, @ToolParam(description = "要加载的类名，支持全限定名", required = false) - String loadClass) { + String loadClass, + + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("classloader"); @@ -68,6 +75,6 @@ public class ClassLoaderTool { cmd.append(" --load ").append(loadClass.trim()); } - return ArthasCommandExecutor.executeCommand(cmd.toString()); + return JsonParser.toJson(commandContext.executeSync(cmd.toString())); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/DumpClassTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/DumpClassTool.java new file mode 100644 index 000000000..449453a39 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/DumpClassTool.java @@ -0,0 +1,68 @@ +package com.taobao.arthas.mcp.server.tool.function.klass100; + +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import java.nio.file.Paths; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; + +public class DumpClassTool { + + /** + * Dump 命令 - 将JVM中已加载的类字节码导出到指定目录 + * 适用于批量下载指定包目录的class字节码 + */ + @Tool( + name = "dump", + description = "将JVM中实际运行的class字节码dump到指定目录，适用于批量下载指定包目录的class字节码" + ) + public String dump( + @ToolParam(description = "类名表达式匹配，如java.lang.String或demo.MathGame") + String classPattern, + + @ToolParam(description = "指定输出目录，默认为arthas-output文件夹", required = false) + String outputDir, + + @ToolParam(description = "ClassLoader的hashcode（16进制），用于指定特定的ClassLoader", required = false) + String classLoaderHashcode, + + @ToolParam(description = "ClassLoader的完整类名，如sun.misc.Launcher$AppClassLoader，可替代hashcode", required = false) + String classLoaderClass, + + @ToolParam(description = "是否包含子类，默认为false", required = false) + Boolean includeInnerClasses, + + @ToolParam(description = "限制dump的类数量，避免输出过多文件", required = false) + Integer limit, + + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + StringBuilder cmd = new StringBuilder("dump"); + + cmd.append(" ").append(classPattern); + + if (outputDir != null && !outputDir.trim().isEmpty()) { + cmd.append(" -d ").append(outputDir.trim()); + } + + if (classLoaderHashcode != null && !classLoaderHashcode.trim().isEmpty()) { + cmd.append(" -c ").append(classLoaderHashcode.trim()); + } else if (classLoaderClass != null && !classLoaderClass.trim().isEmpty()) { + cmd.append(" --classLoaderClass ").append(classLoaderClass.trim()); + } + + if (Boolean.TRUE.equals(includeInnerClasses)) { + cmd.append(" --include-inner-classes"); + } + + if (limit != null && limit > 0) { + cmd.append(" --limit ").append(limit); + } + + return JsonParser.toJson(commandContext.executeSync(cmd.toString())); + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/JadTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/JadTool.java index 43eceaef9..30f86129b 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/JadTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/JadTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class JadTool { @@ -30,7 +34,10 @@ public class JadTool { Boolean useRegex, @ToolParam(description = "指定dump class文件目录，默认会dump到logback.xml中配置的log目录下", required = false) - String dumpDirectory) { + String dumpDirectory, + + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("jad"); @@ -58,6 +65,6 @@ public class JadTool { cmd.append(" -d ").append(dumpDirectory.trim()); } - return ArthasCommandExecutor.executeCommand(cmd.toString()); + return JsonParser.toJson(commandContext.executeSync(cmd.toString())); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/MemoryCompilerTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/MemoryCompilerTool.java index 6a20620fd..af93f1a5a 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/MemoryCompilerTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/MemoryCompilerTool.java @@ -1,11 +1,19 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import java.nio.file.Paths; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class MemoryCompilerTool { + public static final String DEFAULT_DUMP_DIR = Paths.get("arthas-output").toAbsolutePath().toString(); + @Tool( name = "mc", description = "Memory Compiler/内存编译器，编译.java文件生成.class" @@ -20,8 +28,11 @@ public class MemoryCompilerTool { @ToolParam(description = "ClassLoader的完整类名，如sun.misc.Launcher$AppClassLoader，可替代hashcode", required = false) String classLoaderClass, - @ToolParam(description = "指定输出目录", required = false) - String outputDir) { + @ToolParam(description = "指定输出目录，默认为工作目录下arthas-output文件夹", required = false) + String outputDir, + + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("mc"); @@ -35,8 +46,10 @@ public class MemoryCompilerTool { if (outputDir != null && !outputDir.trim().isEmpty()) { cmd.append(" -d ").append(outputDir.trim()); + }else { + cmd.append(" -d ").append(DEFAULT_DUMP_DIR); } - return ArthasCommandExecutor.executeCommand(cmd.toString()); + return JsonParser.toJson(commandContext.executeSync(cmd.toString())); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RedefineTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RedefineTool.java index 2190f8f2e..d42e41825 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RedefineTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RedefineTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class RedefineTool { @@ -18,7 +22,10 @@ public class RedefineTool { String classLoaderHash, @ToolParam(description = "指定执行表达式的ClassLoader的class name，如sun.misc.Launcher$AppClassLoader，可替代hashcode", required = false) - String classLoaderClass) { + String classLoaderClass, + + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("redefine"); @@ -30,6 +37,6 @@ public class RedefineTool { cmd.append(" --classLoaderClass ").append(classLoaderClass.trim()); } - return ArthasCommandExecutor.executeCommand(cmd.toString()); + return JsonParser.toJson(commandContext.executeSync(cmd.toString())); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RetransformTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RetransformTool.java index f279253e3..3f3139bad 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RetransformTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/RetransformTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class RetransformTool { @@ -18,7 +22,10 @@ public class RetransformTool { String classLoaderHash, @ToolParam(description = "ClassLoader的完整类名，如sun.misc.Launcher$AppClassLoader，可替代hashcode", required = false) - String classLoaderClass) { + String classLoaderClass, + + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); StringBuilder cmd = new StringBuilder("retransform"); @@ -30,6 +37,6 @@ public class RetransformTool { cmd.append(" --classLoaderClass ").append(classLoaderClass.trim()); } - return ArthasCommandExecutor.executeCommand(cmd.toString()); + return JsonParser.toJson(commandContext.executeSync(cmd.toString())); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchClassTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchClassTool.java index e0337e0cb..26f595b8a 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchClassTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchClassTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class SearchClassTool { @@ -11,8 +15,10 @@ public class SearchClassTool { description = "获取指定类的详细信息，包括类加载器、方法、字段等信息" ) public String sc( - @ToolParam(description = "要查询的类名，支持全限定名") String className) { - return ArthasCommandExecutor.executeCommand("sc -d " + className); + @ToolParam(description = "要查询的类名，支持全限定名") String className, + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + return JsonParser.toJson(commandContext.executeSync("sc -d " + className)); } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchMethodTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchMethodTool.java index 8021790d6..c2bcb94b1 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchMethodTool.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/klass100/SearchMethodTool.java @@ -1,8 +1,12 @@ package com.taobao.arthas.mcp.server.tool.function.klass100; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; 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; +import com.taobao.arthas.mcp.server.util.JsonParser; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.TOOL_CONTEXT_COMMAND_CONTEXT_KEY; public class SearchMethodTool { @@ -11,8 +15,10 @@ public class SearchMethodTool { description = "获取指定类的所有方法信息" ) public String sm( - @ToolParam(description = "要查询的类名，支持全限定名") String className) { - return ArthasCommandExecutor.executeCommand("sm " + className); + @ToolParam(description = "要查询的类名，支持全限定名") String className, + ToolContext toolContext) { + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + return JsonParser.toJson(commandContext.executeSync("sm " + className)); } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/MonitorTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/MonitorTool.java new file mode 100644 index 000000000..8c6225f19 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/MonitorTool.java @@ -0,0 +1,122 @@ +package com.taobao.arthas.mcp.server.tool.function.monitor200; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class MonitorTool { + + private static final Logger logger = LoggerFactory.getLogger(MonitorTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_REFRESH_INTERVAL_MS = 3000; + public static final int DEFAULT_MAX_MATCH_COUNT = 50; + + /** + * monitor 方法调用监控工具 + * 实时监控指定类的指定方法的调用情况 + * 支持: + * - classPattern: 类名表达式匹配，支持通配符 + * - methodPattern: 方法名表达式匹配，支持通配符 + * - condition: OGNL条件表达式，满足条件的调用才会被监控 + * - intervalMs: 监控统计输出间隔，默认3000ms + * - numberOfExecutions: 监控执行次数，默认3次 + * - regex: 是否开启正则表达式匹配，默认false + * - excludeClassPattern: 排除的类名模式 + * - maxMatch: 最大匹配类数量，防止匹配过多类影响性能，默认50 + */ + @Tool( + name = "monitor", + description = "Monitor 方法调用监控工具: 实时监控指定类的指定方法的调用情况，包括调用次数、成功次数、失败次数、平均RT、失败率等统计信息。对应 Arthas 的 monitor 命令。", + streamable = true + ) + public String monitor( + @ToolParam(description = "类名表达式匹配，支持通配符，如demo.MathGame") + String classPattern, + + @ToolParam(description = "方法名表达式匹配，支持通配符，如primeFactors", required = false) + String methodPattern, + + @ToolParam(description = "OGNL条件表达式，满足条件的调用才会被监控，如params[0]<0", required = false) + String condition, + + @ToolParam(description = "监控统计输出间隔，单位为毫秒，默认 3000ms。用于控制输出频率", required = false) + Integer intervalMs, + + @ToolParam(description = "执行次数限制，默认值为3。达到指定次数后自动停止", required = false) + Integer numberOfExecutions, + + @ToolParam(description = "开启正则表达式匹配，默认为通配符匹配，默认false", required = false) + Boolean regex, + + @ToolParam(description = "最大匹配类数量，防止匹配过多类影响性能，默认50", required = false) + Integer maxMatch, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + int interval = (intervalMs != null && intervalMs > 0) ? intervalMs : DEFAULT_REFRESH_INTERVAL_MS; + int execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + int maxMatchCount = (maxMatch != null && maxMatch > 0) ? maxMatch : DEFAULT_MAX_MATCH_COUNT; + + try { + StringBuilder cmd = new StringBuilder("monitor"); + + // -c 指定统计周期，默认值为 60 秒 + cmd.append(" -c ").append(interval / 1000); + + cmd.append(" -n ").append(execCount); + + cmd.append(" -m ").append(maxMatchCount); + + if (Boolean.TRUE.equals(regex)) { + cmd.append(" -E"); + } + + if (classPattern != null && !classPattern.trim().isEmpty()) { + cmd.append(" ").append(classPattern.trim()); + } else { + throw new IllegalArgumentException("classPattern is required"); + } + + if (methodPattern != null && !methodPattern.trim().isEmpty()) { + cmd.append(" ").append(methodPattern.trim()); + } + + if (condition != null && !condition.trim().isEmpty()) { + cmd.append(" '").append(condition.trim()).append("'"); + } + + String commandStr = cmd.toString(); + logger.info("Starting monitor execution: {}", commandStr); + + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, interval / 10, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("Monitor execution completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("Monitor execution failed due to timeout or error limits exceeded")); + } + + } catch (Exception e) { + logger.error("Error executing monitor command", e); + return JsonParser.toJson(createErrorResponse("Error executing monitor: " + e.getMessage())); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/StackTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/StackTool.java new file mode 100644 index 000000000..df7294911 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/StackTool.java @@ -0,0 +1,106 @@ +package com.taobao.arthas.mcp.server.tool.function.monitor200; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class StackTool { + + private static final Logger logger = LoggerFactory.getLogger(StackTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_POLL_INTERVAL_MS = 50; + + /** + * stack 调用堆栈跟踪工具 + * 输出当前方法被调用的调用路径 + * 支持: + * - classPattern: 类名表达式匹配，支持通配符 + * - methodPattern: 方法名表达式匹配，支持通配符 + * - condition: OGNL条件表达式，满足条件的调用才会被跟踪 + * - numberOfExecutions: 捕获次数限制，达到指定次数后自动停止 + * - regex: 是否开启正则表达式匹配，默认false + * - excludeClassPattern: 排除的类名模式 + */ + @Tool( + name = "stack", + description = "Stack 调用堆栈跟踪工具: 输出当前方法被调用的调用路径，帮助分析方法的调用链路。对应 Arthas 的 stack 命令。", + streamable = true + ) + public String stack( + @ToolParam(description = "类名表达式匹配，支持通配符，如demo.MathGame") + String classPattern, + + @ToolParam(description = "方法名表达式匹配，支持通配符，如primeFactors", required = false) + String methodPattern, + + @ToolParam(description = "OGNL条件表达式，满足条件的调用才会被跟踪，如params[0]<0", required = false) + String condition, + + @ToolParam(description = "捕获次数限制，默认值为3。达到指定次数后自动停止", required = false) + Integer numberOfExecutions, + + @ToolParam(description = "开启正则表达式匹配，默认为通配符匹配，默认false", required = false) + Boolean regex, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + int execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + + try { + StringBuilder cmd = new StringBuilder("stack"); + + cmd.append(" -n ").append(execCount); + + if (Boolean.TRUE.equals(regex)) { + cmd.append(" -E"); + } + + if (classPattern != null && !classPattern.trim().isEmpty()) { + cmd.append(" ").append(classPattern.trim()); + } else { + throw new IllegalArgumentException("classPattern is required"); + } + + if (methodPattern != null && !methodPattern.trim().isEmpty()) { + cmd.append(" ").append(methodPattern.trim()); + } + + if (condition != null && !condition.trim().isEmpty()) { + cmd.append(" '").append(condition.trim()).append("'"); + } + + String commandStr = cmd.toString(); + logger.info("Starting stack execution: {}", commandStr); + + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, DEFAULT_POLL_INTERVAL_MS, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("Stack execution completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("Stack execution failed due to timeout or error limits exceeded")); + } + + } catch (Exception e) { + logger.error("Error executing stack command", e); + return JsonParser.toJson(createErrorResponse("Error executing stack: " + e.getMessage())); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TimeTunnelTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TimeTunnelTool.java new file mode 100644 index 000000000..87bc21d9c --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TimeTunnelTool.java @@ -0,0 +1,258 @@ +package com.taobao.arthas.mcp.server.tool.function.monitor200; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class TimeTunnelTool { + + private static final Logger logger = LoggerFactory.getLogger(TimeTunnelTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_POLL_INTERVAL_MS = 100; + public static final int DEFAULT_MAX_MATCH_COUNT = 50; + + /** + * tt 时空隧道工具 (TimeTunnel) + * 方法执行数据的时空隧道，记录下指定方法每次调用的入参和返回信息，并能对这些不同的时间下调用进行观测 + */ + @Tool( + name = "tt", + description = "TimeTunnel 时空隧道工具: 方法执行数据的时空隧道，记录下指定方法每次调用的入参和返回信息，对应 Arthas 的 tt 命令。支持记录、列表、搜索、查看详情、重放、删除等操作。", + streamable = true + ) + public String timeTunnel( + @ToolParam(description = "操作类型: record/t(记录), list/l(列表), search/s(搜索), info/i(查看详情), replay/p(重放), delete/d(删除), deleteAll/da(删除所有)，默认record") + String action, + + @ToolParam(description = "类名表达式匹配，支持通配符，如demo.MathGame或*Test。record操作时必需", required = false) + String classPattern, + + @ToolParam(description = "方法名表达式匹配，支持通配符，如primeFactors或*method。record操作时必需", required = false) + String methodPattern, + + @ToolParam(description = "OGNL条件表达式，满足条件的调用才会被记录，如params[0]<0或'params.length==1'", required = false) + String condition, + + @ToolParam(description = "记录次数限制，默认值为 3。达到指定次数后自动停止（仅record操作）", required = false) + Integer numberOfExecutions, + + @ToolParam(description = "开启正则表达式匹配，默认为通配符匹配，默认false", required = false) + Boolean regex, + + @ToolParam(description = "指定索引，用于info/replay/delete等操作", required = false) + Integer index, + + @ToolParam(description = "搜索表达式，用于search操作，支持OGNL表达式如'method.name==\"primeFactors\"'", required = false) + String searchExpression, + + @ToolParam(description = "Class最大匹配数量，防止匹配到的Class数量太多导致JVM挂起，默认50", required = false) + Integer maxMatchCount, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + String ttAction = normalizeAction(action); + int execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + int maxMatch = (maxMatchCount != null && maxMatchCount > 0) ? maxMatchCount : DEFAULT_MAX_MATCH_COUNT; + + validateParameters(ttAction, classPattern, methodPattern, index, searchExpression); + + try { + StringBuilder cmd = new StringBuilder("tt"); + + switch (ttAction) { + case "record": + cmd = buildRecordCommand(cmd, classPattern, methodPattern, condition, execCount, maxMatch, regex); + break; + + case "list": + cmd = buildListCommand(cmd, searchExpression); + break; + + case "info": + cmd.append(" -i ").append(index); + break; + + case "search": + cmd.append(" -s '").append(searchExpression.trim()).append("'"); + break; + + case "replay": + cmd.append(" -i ").append(index).append(" -p"); + break; + + case "delete": + cmd.append(" -i ").append(index).append(" -d"); + break; + + case "deleteall": + cmd.append(" --delete-all"); + break; + + default: + throw new IllegalArgumentException("Unsupported action: " + ttAction + + ". Supported actions: record(t), list(l), info(i), search(s), replay(p), delete(d), deleteAll(da)"); + } + + String commandStr = cmd.toString(); + logger.info("Starting TimeTunnel execution: {}", commandStr); + + if ("record".equals(ttAction)) { + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, DEFAULT_POLL_INTERVAL_MS, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("TimeTunnel recording completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("TimeTunnel recording failed due to timeout or error limits exceeded")); + } + } else { + // list/info/search/replay/delete/deleteAll：使用同步执行，直接返回结果 + logger.info("Executing sync tt command: {}", commandStr); + Map result = commandContext.executeSync(commandStr); + return JsonParser.toJson(result); + } + + } catch (Exception e) { + logger.error("Error executing tt command", e); + return JsonParser.toJson(createErrorResponse("Error executing tt: " + e.getMessage())); + } + } + + /** + * 验证参数 + */ + private void validateParameters(String action, String classPattern, String methodPattern, + Integer index, String searchExpression) { + switch (action) { + case "record": + if (classPattern == null || classPattern.trim().isEmpty()) { + throw new IllegalArgumentException("classPattern is required for record operation"); + } + if (methodPattern == null || methodPattern.trim().isEmpty()) { + throw new IllegalArgumentException("methodPattern is required for record operation"); + } + break; + case "info": + case "replay": + if (index == null) { + throw new IllegalArgumentException(action + " operation requires index parameter"); + } + break; + case "search": + if (searchExpression == null || searchExpression.trim().isEmpty()) { + throw new IllegalArgumentException("search operation requires searchExpression parameter"); + } + break; + case "delete": + if (index == null) { + throw new IllegalArgumentException("delete operation requires index parameter"); + } + break; + case "list": + case "deleteall": + break; + default: + throw new IllegalArgumentException("Unsupported action: " + action); + } + } + + /** + * 标准化操作类型 + */ + private String normalizeAction(String action) { + if (action == null || action.trim().isEmpty()) { + return "record"; + } + + String normalizedAction = action.trim().toLowerCase(); + + switch (normalizedAction) { + case "record": + case "r": + case "-t": + case "t": + return "record"; + + case "list": + case "l": + case "-l": + return "list"; + + case "info": + case "i": + case "-i": + return "info"; + + case "search": + case "s": + case "-s": + return "search"; + + case "replay": + case "p": + case "-p": + return "replay"; + + case "delete": + case "d": + case "-d": + return "delete"; + + case "deleteall": + case "da": + case "--delete-all": + return "deleteall"; + + default: + return normalizedAction; + } + } + + private StringBuilder buildRecordCommand(StringBuilder cmd, String classPattern, String methodPattern, + String condition, int execCount, int maxMatch, Boolean regex) { + + cmd.append(" -t"); + + cmd.append(" -n ").append(execCount); + cmd.append(" -m ").append(maxMatch); + + if (Boolean.TRUE.equals(regex)) { + cmd.append(" -E"); + } + + cmd.append(" '").append(classPattern.trim()).append("'"); + cmd.append(" '").append(methodPattern.trim()).append("'"); + + if (condition != null && !condition.trim().isEmpty()) { + cmd.append(" '").append(condition.trim()).append("'"); + } + + return cmd; + } + + private StringBuilder buildListCommand(StringBuilder cmd, String searchExpression) { + cmd.append(" -l"); + if (searchExpression != null && !searchExpression.trim().isEmpty()) { + cmd.append(" '").append(searchExpression.trim()).append("'"); + } + return cmd; + } + +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TraceTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TraceTool.java new file mode 100644 index 000000000..a5bca7694 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/TraceTool.java @@ -0,0 +1,105 @@ +package com.taobao.arthas.mcp.server.tool.function.monitor200; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class TraceTool { + + private static final Logger logger = LoggerFactory.getLogger(TraceTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_POLL_INTERVAL_MS = 100; + public static final int DEFAULT_MAX_MATCH_COUNT = 50; + + /** + * trace 方法内部调用路径跟踪工具 + * 追踪方法内部调用路径，输出每个节点的耗时信息 + */ + @Tool( + name = "trace", + description = "Trace 方法内部调用路径跟踪工具: 追踪方法内部调用路径，输出每个节点的耗时信息，对应 Arthas 的 trace 命令。", + streamable = true + ) + public String trace( + @ToolParam(description = "类名表达式匹配，支持通配符，如demo.MathGame") + String classPattern, + + @ToolParam(description = "方法名表达式匹配，支持通配符，如primeFactors", required = false) + String methodPattern, + + @ToolParam(description = "OGNL条件表达式，包括#cost耗时过滤，如'#cost>100'表示耗时超过100ms", required = false) + String condition, + + @ToolParam(description = "执行次数限制，默认值为3。达到指定次数后自动停止", required = false) + Integer numberOfExecutions, + + @ToolParam(description = "开启正则表达式匹配，默认为通配符匹配，默认false", required = false) + Boolean regex, + + @ToolParam(description = "指定Class最大匹配数量，默认50", required = false) + Integer maxMatchCount, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + int execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + int maxMatch = (maxMatchCount != null && maxMatchCount > 0) ? maxMatchCount : DEFAULT_MAX_MATCH_COUNT; + + try { + StringBuilder cmd = new StringBuilder("trace"); + + cmd.append(" -n ").append(execCount); + cmd.append(" -m ").append(maxMatch); + + if (Boolean.TRUE.equals(regex)) { + cmd.append(" -E"); + } + + if (classPattern != null && !classPattern.trim().isEmpty()) { + cmd.append(" ").append(classPattern.trim()); + } else { + throw new IllegalArgumentException("classPattern is required"); + } + + if (methodPattern != null && !methodPattern.trim().isEmpty()) { + cmd.append(" ").append(methodPattern.trim()); + } + + if (condition != null && !condition.trim().isEmpty()) { + cmd.append(" '").append(condition.trim()).append("'"); + } + + String commandStr = cmd.toString(); + logger.info("Starting trace execution: {}", commandStr); + + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, DEFAULT_POLL_INTERVAL_MS, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("Trace execution completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("Trace execution failed due to timeout or error limits exceeded")); + } + + } catch (Exception e) { + logger.error("Error executing trace command", e); + return JsonParser.toJson(createErrorResponse("Error executing trace: " + e.getMessage())); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/WatchTool.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/WatchTool.java new file mode 100644 index 000000000..4cd274d68 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/function/monitor200/WatchTool.java @@ -0,0 +1,138 @@ +package com.taobao.arthas.mcp.server.tool.function.monitor200; + +import com.taobao.arthas.mcp.server.protocol.server.McpNettyServerExchange; +import com.taobao.arthas.mcp.server.session.ArthasCommandContext; +import com.taobao.arthas.mcp.server.tool.ToolContext; +import com.taobao.arthas.mcp.server.tool.annotation.Tool; +import com.taobao.arthas.mcp.server.tool.annotation.ToolParam; +import com.taobao.arthas.mcp.server.util.JsonParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +import static com.taobao.arthas.mcp.server.tool.util.McpToolUtils.*; +import static com.taobao.arthas.mcp.server.tool.function.StreamableToolUtils.*; + +public class WatchTool { + + private static final Logger logger = LoggerFactory.getLogger(WatchTool.class); + + public static final int DEFAULT_NUMBER_OF_EXECUTIONS = 3; + public static final int DEFAULT_POLL_INTERVAL_MS = 50; + public static final int DEFAULT_MAX_MATCH_COUNT = 50; + public static final int DEFAULT_EXPAND_LEVEL = 1; + public static final String DEFAULT_EXPRESS = "{params, target, returnObj}"; + + /** + * watch 方法执行观察工具 + * 观察指定方法的调用情况，包括入参、返回值和抛出异常等信息 + * 整合了完整的参数支持和流式输出能力 + */ + @Tool( + name = "watch", + description = "Watch 方法执行观察工具: 观察指定方法的调用情况，包括入参、返回值和抛出异常等信息，支持实时流式输出。对应 Arthas 的 watch 命令。", + streamable = true + ) + public String watch( + @ToolParam(description = "类名表达式匹配，支持通配符，如demo.MathGame") + String classPattern, + + @ToolParam(description = "方法名表达式匹配，支持通配符，如primeFactors", required = false) + String methodPattern, + + @ToolParam(description = "观察表达式，默认为{params, target, returnObj}，支持OGNL表达式", required = false) + String express, + + @ToolParam(description = "OGNL条件表达式，满足条件的调用才会被观察，如params[0]<0", required = false) + String condition, + + @ToolParam(description = "在方法调用之前观察(-b)，默认false", required = false) + Boolean beforeMethod, + + @ToolParam(description = "在方法抛出异常后观察(-e)，默认false", required = false) + Boolean exceptionOnly, + + @ToolParam(description = "在方法正常返回后观察(-s)，默认false", required = false) + Boolean successOnly, + + @ToolParam(description = "执行次数限制，默认值为 3。达到指定次数后自动停止", required = false) + Integer numberOfExecutions, + + @ToolParam(description = "开启正则表达式匹配，默认为通配符匹配，默认false", required = false) + Boolean regex, + + @ToolParam(description = "指定Class最大匹配数量，默认50", required = false) + Integer maxMatchCount, + + @ToolParam(description = "指定输出结果的属性遍历深度，默认1，最大4", required = false) + Integer expandLevel, + + ToolContext toolContext + ) { + McpNettyServerExchange exchange = (McpNettyServerExchange) toolContext.getContext().get(TOOL_CONTEXT_MCP_EXCHANGE_KEY); + ArthasCommandContext commandContext = (ArthasCommandContext) toolContext.getContext().get(TOOL_CONTEXT_COMMAND_CONTEXT_KEY); + Object progressTokenObj = toolContext.getContext().get(PROGRESS_TOKEN); + String progressToken = progressTokenObj != null ? String.valueOf(progressTokenObj) : null; + + int execCount = (numberOfExecutions != null && numberOfExecutions > 0) ? numberOfExecutions : DEFAULT_NUMBER_OF_EXECUTIONS; + int maxMatch = (maxMatchCount != null && maxMatchCount > 0) ? maxMatchCount : DEFAULT_MAX_MATCH_COUNT; + int expandDepth = (expandLevel != null && expandLevel >= 1 && expandLevel <= 4) ? expandLevel : DEFAULT_EXPAND_LEVEL; + String watchExpress = (express != null && !express.trim().isEmpty()) ? express.trim() : DEFAULT_EXPRESS; + + try { + StringBuilder cmd = new StringBuilder("watch"); + + cmd.append(" -n ").append(execCount); + cmd.append(" -m ").append(maxMatch); + cmd.append(" -x ").append(expandDepth); + + if (Boolean.TRUE.equals(regex)) { + cmd.append(" -E"); + } + + if (Boolean.TRUE.equals(beforeMethod)) { + cmd.append(" -b"); + } else if (Boolean.TRUE.equals(exceptionOnly)) { + cmd.append(" -e"); + } else if (Boolean.TRUE.equals(successOnly)) { + cmd.append(" -s"); + } else { + cmd.append(" -f"); + } + + if (classPattern != null && !classPattern.trim().isEmpty()) { + cmd.append(" ").append(classPattern.trim()); + } else { + throw new IllegalArgumentException("classPattern is required"); + } + + if (methodPattern != null && !methodPattern.trim().isEmpty()) { + cmd.append(" ").append(methodPattern.trim()); + } + + cmd.append(" '").append(watchExpress).append("'"); + + if (condition != null && !condition.trim().isEmpty()) { + cmd.append(" '").append(condition.trim()).append("'"); + } + + String commandStr = cmd.toString(); + logger.info("Starting watch execution: {}", commandStr); + + Map asyncResult = commandContext.executeAsync(commandStr); + logger.debug("Async execution started: {}", asyncResult); + + Map results = executeAndCollectResults(exchange, commandContext, execCount, DEFAULT_POLL_INTERVAL_MS, progressToken); + if (results != null) { + return JsonParser.toJson(createCompletedResponse("Watch execution completed successfully", results)); + } else { + return JsonParser.toJson(createErrorResponse("Watch execution failed due to timeout or error limits exceeded")); + } + + } catch (Exception e) { + logger.error("Error executing watch command", e); + return JsonParser.toJson(createErrorResponse("Error executing watch: " + e.getMessage())); + } + } +} diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/JsonSchemaGenerator.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/JsonSchemaGenerator.java index 38b87ee5c..d9d9d9702 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/JsonSchemaGenerator.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/JsonSchemaGenerator.java @@ -45,6 +45,11 @@ public final class JsonSchemaGenerator { Parameter[] parameters = method.getParameters(); for (Parameter parameter : parameters) { + ToolParam toolParam = parameter.getAnnotation(ToolParam.class); + if (toolParam == null) { + continue; + } + String paramName = getParameterName(parameter); Class paramType = parameter.getType(); diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/McpToolUtils.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/McpToolUtils.java index dccc0da31..3b480c55a 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/McpToolUtils.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/util/McpToolUtils.java @@ -1,85 +1,148 @@ 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.server.McpStatelessServerFeatures; 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.*; 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"; + public static final String TOOL_CONTEXT_COMMAND_CONTEXT_KEY = "commandContext"; + + public static final String MCP_TRANSPORT_CONTEXT = "mcpTransportContext"; + + public static final String PROGRESS_TOKEN = "progressToken"; + private McpToolUtils() { } - public static List toToolSpecifications( - List tools, - McpServerProperties serverProperties) { + public static List toStreamableToolSpecifications( + List tools) { - // 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()); - } + if (tools == null || tools.isEmpty()) { + return Collections.emptyList(); + } - public static McpServerFeatures.ToolSpecification toToolSpecification(ToolCallback toolCallback) { + // De-duplicate tools by their name, keeping the first occurrence of each tool name + return tools.stream() + .filter(Objects::nonNull) + .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() + .filter(McpToolUtils::isStreamableTool) + .map(McpToolUtils::toToolSpecification) + .collect(Collectors.toList()); + } - McpSchema.Tool tool = new McpSchema.Tool(toolCallback.getToolDefinition().getName(), - toolCallback.getToolDefinition().getDescription(), toolCallback.getToolDefinition().getInputSchema()); + 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, CompletableFuture>() { + McpServerFeatures.ToolCallFunction callFunction = (exchange, commandContext, request) -> { + try { + Map contextMap = new HashMap<>(); + contextMap.put(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchange); + contextMap.put(TOOL_CONTEXT_COMMAND_CONTEXT_KEY, commandContext); + contextMap.put(PROGRESS_TOKEN, request.getMeta().get(PROGRESS_TOKEN)); + ToolContext toolContext = new ToolContext(contextMap); - @Override - public CompletableFuture apply(McpNettyServerExchange exchange, Map request) { - try { - Map contextMap = new HashMap<>(); - contextMap.put(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchange); - ToolContext toolContext = new ToolContext(contextMap); + String requestJson = convertParametersToString(request.getArguments()); - String callResult = toolCallback.call(convertRequestToString(request), toolContext); + String callResult = toolCallback.call(requestJson, toolContext); + return CompletableFuture.completedFuture(createSuccessResult(callResult)); + } catch (Exception e) { + return CompletableFuture.completedFuture(createErrorResult(e.getMessage())); + } + }; + return new McpServerFeatures.ToolSpecification(tool, callFunction); + } - List contents = new ArrayList<>(); - contents.add(new McpSchema.TextContent(callResult)); - return CompletableFuture.completedFuture(new McpSchema.CallToolResult(contents, false)); - } - catch (Exception e) { - List contents = new ArrayList<>(); - contents.add(new McpSchema.TextContent(e.getMessage())); + public static List toStatelessToolSpecifications(List providerToolCallbacks) { + if (providerToolCallbacks == null || providerToolCallbacks.isEmpty()) { + return Collections.emptyList(); + } - return CompletableFuture.completedFuture(new McpSchema.CallToolResult(contents, true)); - } - } + // De-duplicate tools by their name, keeping the first occurrence of each tool name + return providerToolCallbacks.stream() + .filter(Objects::nonNull) + .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::toStatelessToolSpecification) + .collect(Collectors.toList()); + } - private String convertRequestToString(Map request) { - if (request == null) { - return ""; - } - try { - return new ObjectMapper().writeValueAsString(request); - } catch (Exception e) { - return request.toString(); - } - } - }); - } + public static McpStatelessServerFeatures.ToolSpecification toStatelessToolSpecification(ToolCallback toolCallback) { + McpSchema.Tool tool = new McpSchema.Tool( + toolCallback.getToolDefinition().getName(), + toolCallback.getToolDefinition().getDescription(), + toolCallback.getToolDefinition().getInputSchema() + ); + + McpStatelessServerFeatures.ToolCallFunction callFunction = (context, commandContext, arguments) -> { + try { + Map contextMap = new HashMap<>(); + contextMap.put(MCP_TRANSPORT_CONTEXT, context); + contextMap.put(TOOL_CONTEXT_COMMAND_CONTEXT_KEY, commandContext); + ToolContext toolContext = new ToolContext(contextMap); + + String argumentsJson = convertParametersToString(arguments); + String callResult = toolCallback.call(argumentsJson, toolContext); + return CompletableFuture.completedFuture(createSuccessResult(callResult)); + } catch (Exception e) { + return CompletableFuture.completedFuture(createErrorResult("Error executing tool: " + e.getMessage())); + } + }; + + return new McpStatelessServerFeatures.ToolSpecification(tool, callFunction); + } + + public static boolean isStreamableTool(ToolCallback toolCallback) { + return toolCallback.getToolDefinition().isStreamable(); + } + + + private static String convertParametersToString(Map parameters) { + if (parameters == null) { + return ""; + } + try { + return new ObjectMapper().writeValueAsString(parameters); + } catch (Exception e) { + return parameters.toString(); + } + } + + private static McpSchema.CallToolResult createSuccessResult(String content) { + List contents = new ArrayList<>(); + contents.add(new McpSchema.TextContent(content)); + return new McpSchema.CallToolResult(contents, false); + } + + private static McpSchema.CallToolResult createErrorResult(String errorMessage) { + List contents = new ArrayList<>(); + contents.add(new McpSchema.TextContent(errorMessage)); + return new McpSchema.CallToolResult(contents, true); + } } diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/JsonParser.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/JsonParser.java index 7794228ea..f75b41b1b 100644 --- a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/JsonParser.java +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/JsonParser.java @@ -6,6 +6,7 @@ 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 com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; import java.lang.reflect.Type; import java.math.BigDecimal; @@ -18,6 +19,8 @@ public final class JsonParser { private static final ObjectMapper OBJECT_MAPPER = JsonMapper.builder() .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES) .disable(SerializationFeature.FAIL_ON_EMPTY_BEANS) + .addModule(new JavaTimeModule()) + .disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS) .build(); private JsonParser() { diff --git a/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/KeepAliveScheduler.java b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/KeepAliveScheduler.java new file mode 100644 index 000000000..252930e04 --- /dev/null +++ b/labs/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/util/KeepAliveScheduler.java @@ -0,0 +1,209 @@ +/** + * Copyright 2025 - 2025 the original author or authors. + */ + +package com.taobao.arthas.mcp.server.util; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.taobao.arthas.mcp.server.protocol.spec.McpSchema; +import com.taobao.arthas.mcp.server.protocol.spec.McpSession; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; +import java.util.Collection; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; + +/** + * A utility class for scheduling regular keep-alive calls to maintain connections. It + * sends periodic keep-alive, ping, messages to connected mcp clients to prevent idle + * timeouts. + * + * The pings are sent to all active mcp sessions at regular intervals. + * + */ +public class KeepAliveScheduler { + + private static final Logger logger = LoggerFactory.getLogger(KeepAliveScheduler.class); + + private static final TypeReference OBJECT_TYPE_REF = new TypeReference() { + }; + + /** Initial delay before the first keepAlive call */ + private final Duration initialDelay; + + /** Interval between subsequent keepAlive calls */ + private final Duration interval; + + /** The scheduler used for executing keepAlive calls */ + private final ScheduledExecutorService scheduler; + + /** Whether this scheduler owns the executor and should shut it down */ + private final boolean ownsExecutor; + + /** The current state of the scheduler */ + private final AtomicBoolean isRunning = new AtomicBoolean(false); + + /** The current scheduled task */ + private volatile ScheduledFuture currentTask; + + /** Supplier for McpSession instances */ + private final Supplier> mcpSessions; + + /** + * Creates a KeepAliveScheduler with a custom scheduler, initial delay, interval and a + * supplier for McpSession instances. + * @param scheduler The scheduler to use for executing keepAlive calls + * @param ownsExecutor Whether this scheduler owns the executor and should shut it down + * @param initialDelay Initial delay before the first keepAlive call + * @param interval Interval between subsequent keepAlive calls + * @param mcpSessions Supplier for McpSession instances + */ + private KeepAliveScheduler(ScheduledExecutorService scheduler, boolean ownsExecutor, Duration initialDelay, + Duration interval, Supplier> mcpSessions) { + this.scheduler = scheduler; + this.ownsExecutor = ownsExecutor; + this.initialDelay = initialDelay; + this.interval = interval; + this.mcpSessions = mcpSessions; + } + + public static Builder builder(Supplier> mcpSessions) { + return new Builder(mcpSessions); + } + + /** + * Starts regular keepAlive calls with sessions supplier. + * @return This scheduler instance for method chaining + */ + public KeepAliveScheduler start() { + if (this.isRunning.compareAndSet(false, true)) { + logger.debug("Starting KeepAlive scheduler with initial delay: {}ms, interval: {}ms", + initialDelay.toMillis(), interval.toMillis()); + + this.currentTask = this.scheduler.scheduleAtFixedRate( + this::sendKeepAlivePings, + this.initialDelay.toMillis(), + this.interval.toMillis(), + TimeUnit.MILLISECONDS + ); + + return this; + } else { + throw new IllegalStateException("KeepAlive scheduler is already running. Stop it first."); + } + } + + /** + * Sends keep-alive pings to all active sessions. + */ + private void sendKeepAlivePings() { + try { + Collection sessions = this.mcpSessions.get(); + if (sessions == null || sessions.isEmpty()) { + logger.trace("No active sessions to ping"); + return; + } + + logger.trace("Sending keep-alive pings to {} sessions", sessions.size()); + + for (McpSession session : sessions) { + try { + session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF) + .whenComplete((result, error) -> { + if (error != null) { + logger.warn("Failed to send keep-alive ping to session {}: {}", + session, error.getMessage()); + } else { + logger.trace("Keep-alive ping sent successfully to session {}", session); + } + }); + } catch (Exception e) { + logger.warn("Exception while sending keep-alive ping to session {}: {}", + session, e.getMessage()); + } + } + } catch (Exception e) { + logger.error("Error during keep-alive ping cycle", e); + } + } + + public void stop() { + if (this.currentTask != null && !this.currentTask.isCancelled()) { + this.currentTask.cancel(false); + logger.debug("KeepAlive scheduler stopped"); + } + this.isRunning.set(false); + } + + public boolean isRunning() { + return this.isRunning.get(); + } + + public void shutdown() { + stop(); + if (this.ownsExecutor && !this.scheduler.isShutdown()) { + this.scheduler.shutdown(); + try { + if (!this.scheduler.awaitTermination(5, TimeUnit.SECONDS)) { + this.scheduler.shutdownNow(); + } + } catch (InterruptedException e) { + this.scheduler.shutdownNow(); + Thread.currentThread().interrupt(); + } + logger.debug("KeepAlive scheduler executor shut down"); + } + } + + public static class Builder { + + private ScheduledExecutorService scheduler; + private boolean ownsExecutor = false; + private Duration initialDelay = Duration.ofSeconds(0); + private Duration interval = Duration.ofSeconds(30); + private Supplier> mcpSessions; + + Builder(Supplier> mcpSessions) { + Assert.notNull(mcpSessions, "McpSessions supplier must not be null"); + this.mcpSessions = mcpSessions; + } + + public Builder scheduler(ScheduledExecutorService scheduler) { + Assert.notNull(scheduler, "Scheduler must not be null"); + this.scheduler = scheduler; + this.ownsExecutor = false; + return this; + } + + public Builder initialDelay(Duration initialDelay) { + Assert.notNull(initialDelay, "Initial delay must not be null"); + this.initialDelay = initialDelay; + return this; + } + + public Builder interval(Duration interval) { + Assert.notNull(interval, "Interval must not be null"); + this.interval = interval; + return this; + } + + public KeepAliveScheduler build() { + if (this.scheduler == null) { + this.scheduler = Executors.newSingleThreadScheduledExecutor(r -> { + Thread t = new Thread(r, "mcp-keep-alive-scheduler"); + t.setDaemon(true); + return t; + }); + this.ownsExecutor = true; + } + + return new KeepAliveScheduler(scheduler, ownsExecutor, initialDelay, interval, mcpSessions); + } + } +} diff --git a/site/docs/doc/monitor.md b/site/docs/doc/monitor.md index 0d71db688..56de016a9 100644 --- a/site/docs/doc/monitor.md +++ b/site/docs/doc/monitor.md @@ -37,7 +37,7 @@ | _method-pattern_ | 方法名表达式匹配 | | _condition-express_ | 条件表达式 | | [E] | 开启正则表达式匹配，默认为通配符匹配 | -| `[c:]` | 统计周期，默认值为 120 秒 | +| `[c:]` | 统计周期，默认值为 60 秒 | | [b] | 在**方法调用之前**计算 condition-express | | `[m ]` | 指定 Class 最大匹配数量，默认值为 50。长格式为`[maxMatch ]`。 | diff --git a/site/docs/en/doc/monitor.md b/site/docs/en/doc/monitor.md index 9d00ec41c..534fa743b 100644 --- a/site/docs/en/doc/monitor.md +++ b/site/docs/en/doc/monitor.md @@ -37,7 +37,7 @@ Parameter `[c:]` stands for cycles of statistics. Its value is an integer value | _method-pattern_ | pattern for the method name | | _condition-expression_ | condition expression for filtering method calls | | `[E]` | turn on regex matching while the default is wildcard matching | -| `[c:]` | cycle of statistics, the default value: `120`s | +| `[c:]` | cycle of statistics, the default value: `60`s | | `[b]` | evaluate the condition-expression before method invoke | | `[m ]` | Specify the max number of matched Classes, the default value is 50. Long format is `[maxMatch ]`. |