4
4
import com .taobao .arthas .mcp .server .protocol .config .McpServerProperties ;
5
5
import com .taobao .arthas .mcp .server .protocol .server .McpNettyServer ;
6
6
import com .taobao .arthas .mcp .server .protocol .server .McpServer ;
7
- import com .taobao .arthas .mcp .server .protocol .server .handler .McpRequestHandler ;
8
- import com .taobao .arthas .mcp .server .protocol .server .transport .HttpNettyServerTransportProvider ;
9
- import com .taobao .arthas .mcp .server .protocol .spec .McpSchema .*;
10
- import com .taobao .arthas .mcp .server .protocol .spec .McpServerTransportProvider ;
7
+ import com .taobao .arthas .mcp .server .protocol .server .McpStatelessNettyServer ;
8
+ import com .taobao .arthas .mcp .server .protocol .server .handler .McpHttpRequestHandler ;
9
+ import com .taobao .arthas .mcp .server .protocol .server .handler .McpStatelessHttpRequestHandler ;
10
+ import com .taobao .arthas .mcp .server .protocol .server .handler .McpStreamableHttpRequestHandler ;
11
+ import com .taobao .arthas .mcp .server .protocol .server .transport .NettyStatelessServerTransport ;
12
+ import com .taobao .arthas .mcp .server .protocol .server .transport .NettyStreamableServerTransportProvider ;
13
+ import com .taobao .arthas .mcp .server .protocol .spec .McpSchema .Implementation ;
14
+ import com .taobao .arthas .mcp .server .protocol .spec .McpSchema .ServerCapabilities ;
15
+ import com .taobao .arthas .mcp .server .protocol .spec .McpStreamableServerTransportProvider ;
11
16
import com .taobao .arthas .mcp .server .tool .DefaultToolCallbackProvider ;
12
17
import com .taobao .arthas .mcp .server .tool .ToolCallback ;
13
18
import com .taobao .arthas .mcp .server .tool .ToolCallbackProvider ;
14
19
import com .taobao .arthas .mcp .server .tool .util .McpToolUtils ;
20
+ import com .taobao .arthas .mcp .server .util .JsonParser ;
15
21
import org .slf4j .Logger ;
16
22
import org .slf4j .LoggerFactory ;
17
23
18
- import java .util .Arrays ;
19
- import java .util .List ;
20
- import java .util .Objects ;
24
+ import java .util .*;
21
25
import java .util .stream .Collectors ;
22
26
23
27
/**
26
30
*/
27
31
public class ArthasMcpServer {
28
32
private static final Logger logger = LoggerFactory .getLogger (ArthasMcpServer .class );
29
- private McpNettyServer server ;
33
+
34
+ private McpNettyServer streamableServer ;
35
+ private McpStatelessNettyServer statelessServer ;
36
+
30
37
private final int port ;
31
38
private final String bindAddress ;
32
- private McpRequestHandler mcpRequestHandler ;
39
+
40
+ private final CommandExecutor commandExecutor ;
41
+
42
+ private McpHttpRequestHandler unifiedMcpHandler ;
43
+
44
+ private McpStreamableHttpRequestHandler streamableHandler ;
45
+
46
+ private McpStatelessHttpRequestHandler statelessHandler ;
33
47
34
- public ArthasMcpServer () {
35
- this (8080 , "localhost" );
48
+ public ArthasMcpServer (CommandExecutor commandExecutor ) {
49
+ this (8080 , "localhost" , commandExecutor );
36
50
}
37
51
38
- public ArthasMcpServer (int port , String bindAddress ) {
52
+ public ArthasMcpServer (int port , String bindAddress , CommandExecutor commandExecutor ) {
39
53
this .port = port ;
40
54
this .bindAddress = bindAddress ;
55
+ this .commandExecutor = commandExecutor ;
41
56
}
42
57
43
- public McpRequestHandler getMcpRequestHandler () {
44
- return mcpRequestHandler ;
58
+ public McpHttpRequestHandler getMcpRequestHandler () {
59
+ return unifiedMcpHandler ;
45
60
}
46
61
47
62
/**
48
63
* Start MCP server
49
64
*/
50
65
public void start () {
66
+ logger .info ("Starting Arthas MCP server on {}:{}" , bindAddress , port );
51
67
try {
52
- // 1. Create server configuration
53
68
McpServerProperties properties = new McpServerProperties .Builder ()
54
69
.name ("arthas-mcp-server" )
55
70
.version ("1.0.0" )
56
71
.bindAddress (bindAddress )
57
72
.port (port )
58
- .messageEndpoint ("/sse/message" )
59
- .sseEndpoint ("/sse" )
73
+ .mcpEndpoint ("/mcp" )
60
74
.toolChangeNotification (true )
61
75
.resourceChangeNotification (true )
62
76
.promptChangeNotification (true )
63
- .objectMapper (new ObjectMapper ())
77
+ .objectMapper (JsonParser . getObjectMapper ())
64
78
.build ();
65
-
66
- // 2. Create transport provider
67
- McpServerTransportProvider transportProvider = createHttpTransportProvider (properties );
68
- mcpRequestHandler = transportProvider .getMcpRequestHandler ();
69
-
70
- // 3. Create server builder
71
- McpServer .NettySpecification serverBuilder = McpServer .netty (transportProvider )
72
- .serverInfo (new Implementation (properties .getName (), properties .getVersion ()))
73
- .capabilities (buildServerCapabilities (properties ))
74
- .instructions (properties .getInstructions ())
75
- .requestTimeout (properties .getRequestTimeout ())
76
- .objectMapper (properties .getObjectMapper () != null ? properties .getObjectMapper () : new ObjectMapper ());
77
79
78
80
ToolCallbackProvider toolCallbackProvider = new DefaultToolCallbackProvider ();
79
81
ToolCallback [] callbacks = toolCallbackProvider .getToolCallbacks ();
80
82
List <ToolCallback > providerToolCallbacks = Arrays .stream (callbacks )
81
83
.filter (Objects ::nonNull )
82
84
.collect (Collectors .toList ());
83
85
84
- serverBuilder .tools (
85
- McpToolUtils .toToolSpecifications (providerToolCallbacks , properties ));
86
- server = serverBuilder .build ();
86
+ // Create transport for both streamable and stateless servers
87
+ McpStreamableServerTransportProvider transportProvider = createStreamableHttpTransportProvider (properties );
88
+ streamableHandler = transportProvider .getMcpRequestHandler ();
89
+
90
+ NettyStatelessServerTransport statelessTransport = createStatelessHttpTransport (properties );
91
+ statelessHandler = statelessTransport .getMcpRequestHandler ();
92
+
93
+ unifiedMcpHandler = McpHttpRequestHandler .builder ()
94
+ .mcpEndpoint (properties .getMcpEndpoint ())
95
+ .objectMapper (properties .getObjectMapper ())
96
+ .tools (Arrays .asList (callbacks ))
97
+ .build ();
98
+ unifiedMcpHandler .setStreamableHandler (streamableHandler );
99
+ unifiedMcpHandler .setStatelessHandler (statelessHandler );
100
+
101
+ // Set up unified MCP handler for both streamable and stateless servers
102
+ McpServer .StreamableServerNettySpecification streamableServerNettySpecification = McpServer .netty (transportProvider )
103
+ .serverInfo (new Implementation (properties .getName (), properties .getVersion ()))
104
+ .capabilities (buildServerCapabilities (properties ))
105
+ .instructions (properties .getInstructions ())
106
+ .requestTimeout (properties .getRequestTimeout ())
107
+ .commandExecutor (commandExecutor )
108
+ .objectMapper (properties .getObjectMapper () != null ? properties .getObjectMapper () : JsonParser .getObjectMapper ());
109
+
110
+ // Set up unified MCP handler for both streamable and stateless servers
111
+ McpServer .StatelessServerNettySpecification statelessServerNettySpecification = McpServer .netty (statelessTransport )
112
+ .serverInfo (new Implementation (properties .getName (), properties .getVersion ()))
113
+ .capabilities (buildServerCapabilities (properties ))
114
+ .instructions (properties .getInstructions ())
115
+ .requestTimeout (properties .getRequestTimeout ())
116
+ .commandExecutor (commandExecutor )
117
+ .objectMapper (properties .getObjectMapper () != null ? properties .getObjectMapper () : JsonParser .getObjectMapper ());
118
+
119
+ streamableServerNettySpecification .tools (
120
+ McpToolUtils .toStreamableToolSpecifications (providerToolCallbacks ));
121
+ statelessServerNettySpecification .tools (
122
+ McpToolUtils .toStatelessToolSpecifications (providerToolCallbacks ));
123
+
124
+ streamableServer = streamableServerNettySpecification .build ();
125
+ statelessServer = statelessServerNettySpecification .build ();
87
126
88
- logger .info ("Arthas MCP server started, listening on {}:{}" , bindAddress , port );
127
+ logger .info ("Arthas MCP server started successfully" );
128
+ logger .info ("- Listening on {}:{}" , bindAddress , port );
129
+ logger .info ("- MCP Endpoint: {}" , properties .getMcpEndpoint ());
130
+ logger .info ("- Transport modes: Streamable + Stateless" );
131
+ logger .info ("- Available tools: {}" , providerToolCallbacks .size ());
132
+ logger .info ("- Server ready to accept connections" );
89
133
} catch (Exception e ) {
90
134
logger .error ("Failed to start Arthas MCP server" , e );
91
135
throw new RuntimeException ("Failed to start Arthas MCP server" , e );
@@ -95,47 +139,52 @@ public void start() {
95
139
/**
96
140
* Create HTTP transport provider
97
141
*/
98
- private HttpNettyServerTransportProvider createHttpTransportProvider (McpServerProperties properties ) {
99
- return HttpNettyServerTransportProvider .builder ()
100
- .messageEndpoint (properties .getMessageEndpoint ())
101
- .sseEndpoint (properties .getSseEndpoint ())
142
+ private NettyStreamableServerTransportProvider createStreamableHttpTransportProvider (McpServerProperties properties ) {
143
+ return NettyStreamableServerTransportProvider .builder ()
144
+ .mcpEndpoint (properties .getMcpEndpoint ())
145
+ .objectMapper (properties .getObjectMapper () != null ? properties .getObjectMapper () : new ObjectMapper ())
146
+ .build ();
147
+ }
148
+
149
+ private NettyStatelessServerTransport createStatelessHttpTransport (McpServerProperties properties ) {
150
+ return NettyStatelessServerTransport .builder ()
151
+ .mcpEndpoint (properties .getMcpEndpoint ())
102
152
.objectMapper (properties .getObjectMapper () != null ? properties .getObjectMapper () : new ObjectMapper ())
103
153
.build ();
104
154
}
105
155
106
- /**
107
- * Build server capabilities configuration
108
- */
109
156
private ServerCapabilities buildServerCapabilities (McpServerProperties properties ) {
110
157
return ServerCapabilities .builder ()
111
158
.prompts (new ServerCapabilities .PromptCapabilities (properties .isPromptChangeNotification ()))
112
159
.resources (new ServerCapabilities .ResourceCapabilities (properties .isResourceSubscribe (), properties .isResourceChangeNotification ()))
113
160
.tools (new ServerCapabilities .ToolCapabilities (properties .isToolChangeNotification ()))
114
161
.build ();
115
162
}
116
-
117
- /**
118
- * Stop MCP server
119
- */
163
+
120
164
public void stop () {
121
- if ( server != null ) {
122
- try {
123
- server . closeGracefully (). get ();
124
- logger .info ( "Arthas MCP server stopped " );
125
- } catch ( Exception e ) {
126
- logger .error ( "Failed to stop Arthas MCP server" , e );
165
+ logger . info ( "Stopping Arthas MCP server..." );
166
+ try {
167
+ if ( unifiedMcpHandler != null ) {
168
+ logger .debug ( "Shutting down unified MCP handler " );
169
+ unifiedMcpHandler . closeGracefully (). get ();
170
+ logger .info ( "Unified MCP handler stopped successfully" );
127
171
}
128
- }
129
- }
130
172
131
- public static void main (String [] args ) {
132
- ArthasMcpServer arthasMcpServer = new ArthasMcpServer ();
133
- arthasMcpServer .start ();
134
- // Keep the server running
135
- try {
136
- Thread .sleep (Long .MAX_VALUE );
137
- } catch (InterruptedException e ) {
138
- Thread .currentThread ().interrupt ();
173
+ if (streamableServer != null ) {
174
+ logger .debug ("Shutting down streamable server" );
175
+ streamableServer .closeGracefully ().get ();
176
+ logger .info ("Streamable server stopped successfully" );
177
+ }
178
+
179
+ if (statelessServer != null ) {
180
+ logger .debug ("Shutting down stateless server" );
181
+ statelessServer .closeGracefully ().get ();
182
+ logger .info ("Stateless server stopped successfully" );
183
+ }
184
+
185
+ logger .info ("Arthas MCP server stopped completely" );
186
+ } catch (Exception e ) {
187
+ logger .error ("Failed to stop Arthas MCP server gracefully" , e );
139
188
}
140
189
}
141
190
}
0 commit comments