modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

efactor(mcp): reorganize McpAsyncServer and enhance error handling

Browse source 
- Reorganize McpAsyncServer code into logical sections (Lifecycle, Tool, Resource, Prompt Management)
- (bug) Fix resource registration to use URI instead of name as map key
- Improve error logging in StdioServerTransport with helper methods
- Clean up documentation and add transport mode descriptions
- Update sample server configuration with clearer method names and organization
This commit is contained in:
Christian Tzolov 2024-12-30 16:09:12 +01:00
commit a5394fffd2
5 changed files with 195 additions and 180 deletions
Download patch file Download diff file Expand all files Collapse all files

284
mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java
View file

@ -129,6 +129,52 @@ public class McpAsyncServer {
notificationHandlers);
}
// ---------------------------------------
// Lifecycle Management
// ---------------------------------------
private DefaultMcpSession.RequestHandler initializeRequestHandler() {
return params -> {
McpSchema.InitializeRequest initializeRequest = transport.unmarshalFrom(params,
new TypeReference<McpSchema.InitializeRequest>() {
});
logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}",
initializeRequest.protocolVersion(), initializeRequest.capabilities(),
initializeRequest.clientInfo());
if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(initializeRequest.protocolVersion())) {
return Mono
.<Object>error(new McpError(
"Unsupported protocol version from client: " + initializeRequest.protocolVersion()))
.publishOn(Schedulers.boundedElastic());
}
return Mono
.<Object>just(new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION, this.serverCapabilities,
this.serverInfo, null))
.publishOn(Schedulers.boundedElastic());
};
}
/**
* Gracefully closes the server, allowing any in-progress operations to complete.
* @return A Mono that completes when the server has been closed
*/
public Mono<Void> closeGracefully() {
return this.mcpSession.closeGracefully();
}
/**
* Close the server immediately.
*/
public void close() {
this.mcpSession.close();
}
// ---------------------------------------
// Tool Management
// ---------------------------------------
/**
* Add a new tool registration at runtime.
* @param toolRegistration The tool registration to add
@ -185,6 +231,53 @@ public class McpAsyncServer {
return Mono.error(new McpError("Tool with name '" + toolName + "' not found"));
}
/**
* Notifies clients that the list of available tools has changed.
* @return A Mono that completes when all clients have been notified
*/
public Mono<Void> notifyToolsListChanged() {
return this.mcpSession.sendNotification("notifications/tools/list_changed", null);
}
private DefaultMcpSession.RequestHandler toolsListRequestHandler() {
return params -> {
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.PaginatedRequest>() {
});
List<Tool> tools = this.tools.stream().map(toolRegistration -> {
return toolRegistration.tool();
}).toList();
logger.info("Client tools list request - Cursor: {}", request.cursor());
return Mono.just(new McpSchema.ListToolsResult(tools, null));
};
}
private DefaultMcpSession.RequestHandler toolsCallRequestHandler() {
return params -> {
McpSchema.CallToolRequest callToolRequest = transport.unmarshalFrom(params,
new TypeReference<McpSchema.CallToolRequest>() {
});
Optional<ToolRegistration> toolRegistration = this.tools.stream()
.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
.findAny();
if (toolRegistration.isEmpty()) {
return Mono.<Object>error(new McpError("Tool not found: " + callToolRequest.name()));
}
CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments());
return Mono.just(callResponse);
};
}
// ---------------------------------------
// Resource Management
// ---------------------------------------
/**
* Add a new resource handler at runtime.
* @param resourceHandler The resource handler to add
@ -236,6 +329,43 @@ public class McpAsyncServer {
return Mono.error(new McpError("Resource with URI '" + resourceUri + "' not found"));
}
/**
* Notifies clients that the list of available resources has changed.
* @return A Mono that completes when all clients have been notified
*/
public Mono<Void> notifyResourcesListChanged() {
return this.mcpSession.sendNotification("notifications/resources/list_changed", null);
}
private DefaultMcpSession.RequestHandler resourcesListRequestHandler() {
return params -> {
var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList();
return Mono.just(new McpSchema.ListResourcesResult(resourceList, null));
};
}
private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() {
return params -> Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
}
private DefaultMcpSession.RequestHandler resourcesReadRequestHandler() {
return params -> {
McpSchema.ReadResourceRequest resourceRequest = transport.unmarshalFrom(params,
new TypeReference<McpSchema.ReadResourceRequest>() {
});
var resourceUri = resourceRequest.uri();
if (this.resources.containsKey(resourceUri)) {
return Mono.just(this.resources.get(resourceUri).readHandler().apply(resourceRequest));
}
return Mono.error(new McpError("Resource not found: " + resourceUri));
};
}
// ---------------------------------------
// Prompt Management
// ---------------------------------------
/**
* Add a new prompt handler at runtime.
* @param promptRegistration The prompt handler to add
@ -255,7 +385,11 @@ public class McpAsyncServer {
}
this.prompts.put(promptRegistration.propmpt().name(), promptRegistration);
logger.info("Added prompt handler: {}", promptRegistration.propmpt().name());
// Servers that declared the listChanged capability SHOULD send a notification,
// when the list of available prompts changes
if (this.serverCapabilities.prompts().listChanged()) {
return notifyPromptsListChanged();
}
@ -276,8 +410,11 @@ public class McpAsyncServer {
}
PromptRegistration removed = this.prompts.remove(promptName);
if (removed != null) {
logger.info("Removed prompt handler: {}", promptName);
// Servers that declared the listChanged capability SHOULD send a
// notification, when the list of available prompts changes
if (this.serverCapabilities.prompts().listChanged()) {
return this.notifyPromptsListChanged();
}
@ -286,100 +423,12 @@ public class McpAsyncServer {
return Mono.error(new McpError("Prompt with name '" + promptName + "' not found"));
}
// ---------------------------------------
// Request Handlers
// ---------------------------------------
private DefaultMcpSession.RequestHandler initializeRequestHandler() {
return params -> {
McpSchema.InitializeRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.InitializeRequest>() {
});
logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}",
request.protocolVersion(), request.capabilities(), request.clientInfo());
if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(request.protocolVersion())) {
return Mono
.<Object>error(
new McpError("Unsupported protocol version from client: " + request.protocolVersion()))
.publishOn(Schedulers.boundedElastic());
}
return Mono
.<Object>just(new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION, this.serverCapabilities,
this.serverInfo, null))
.publishOn(Schedulers.boundedElastic());
};
}
private DefaultMcpSession.RequestHandler toolsListRequestHandler() {
return params -> {
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.PaginatedRequest>() {
});
List<Tool> tools = this.tools.stream().map(toolRegistration -> {
return toolRegistration.tool();
}).toList();
logger.info("Client tools list request - Cursor: {}", request.cursor());
return Mono.just(new McpSchema.ListToolsResult(tools, null));
};
}
private DefaultMcpSession.RequestHandler toolsCallRequestHandler() {
return params -> {
McpSchema.CallToolRequest callToolRequest = transport.unmarshalFrom(params,
new TypeReference<McpSchema.CallToolRequest>() {
});
Optional<ToolRegistration> toolRegistration = this.tools.stream()
.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
.findAny();
if (toolRegistration.isEmpty()) {
return Mono.<Object>error(new McpError("Tool not found: " + callToolRequest.name()));
}
CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments());
return Mono.just(callResponse);
};
}
private DefaultMcpSession.RequestHandler resourcesListRequestHandler() {
return params -> {
// McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
// new TypeReference<McpSchema.PaginatedRequest>() {
// });
var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList();
return Mono.just(new McpSchema.ListResourcesResult(resourceList, null));
};
}
private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() {
return params -> {
// McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
// new TypeReference<McpSchema.PaginatedRequest>() {
// });
return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
};
}
private DefaultMcpSession.RequestHandler resourcesReadRequestHandler() {
return params -> {
McpSchema.ReadResourceRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.ReadResourceRequest>() {
});
var resourceUri = request.uri();
if (this.resources.containsKey(resourceUri)) {
return Mono.just(this.resources.get(resourceUri).readHandler().apply(request));
}
return Mono.error(new McpError("Resource not found: " + resourceUri));
};
/**
* Notifies clients that the list of available prompts has changed.
* @return A Mono that completes when all clients have been notified
*/
public Mono<Void> notifyPromptsListChanged() {
return this.mcpSession.sendNotification("notifications/prompts/list_changed", null);
}
private DefaultMcpSession.RequestHandler promptsListRequestHandler() {
@ -396,56 +445,17 @@ public class McpAsyncServer {
private DefaultMcpSession.RequestHandler promptsGetRequestHandler() {
return params -> {
McpSchema.GetPromptRequest request = transport.unmarshalFrom(params,
McpSchema.GetPromptRequest promptRequest = transport.unmarshalFrom(params,
new TypeReference<McpSchema.GetPromptRequest>() {
});
// Implement prompt retrieval logic here
if (this.prompts.containsKey(request.name())) {
return Mono.just(this.prompts.get(request.name()).promptHandler().apply(request));
if (this.prompts.containsKey(promptRequest.name())) {
return Mono.just(this.prompts.get(promptRequest.name()).promptHandler().apply(promptRequest));
}
return Mono.error(new McpError("Prompt not found: " + request.name()));
return Mono.error(new McpError("Prompt not found: " + promptRequest.name()));
};
}
/**
* Notifies clients that the list of available tools has changed.
* @return A Mono that completes when all clients have been notified
*/
public Mono<Void> notifyToolsListChanged() {
return this.mcpSession.sendNotification("notifications/tools/list_changed", null);
}
/**
* Notifies clients that the list of available resources has changed.
* @return A Mono that completes when all clients have been notified
*/
public Mono<Void> notifyResourcesListChanged() {
return this.mcpSession.sendNotification("notifications/resources/list_changed", null);
}
/**
* Notifies clients that the list of available prompts has changed.
* @return A Mono that completes when all clients have been notified
*/
public Mono<Void> notifyPromptsListChanged() {
return this.mcpSession.sendNotification("notifications/prompts/list_changed", null);
}
/**
* Gracefully closes the server, allowing any in-progress operations to complete.
* @return A Mono that completes when the server has been closed
*/
public Mono<Void> closeGracefully() {
return this.mcpSession.closeGracefully();
}
/**
* Close the server immediately.
*/
public void close() {
this.mcpSession.close();
}
}

4
mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java
View file

@ -392,7 +392,7 @@ public interface McpServer {
public Builder resources(List<ResourceRegistration> resourceRegsitrations) {
Assert.notNull(resourceRegsitrations, "Resource handlers list must not be null");
for (ResourceRegistration resource : resourceRegsitrations) {
this.resources.put(resource.resource().name(), resource);
this.resources.put(resource.resource().uri(), resource);
}
return this;
}
@ -417,7 +417,7 @@ public interface McpServer {
public Builder resources(ResourceRegistration... resourceRegistrations) {
Assert.notNull(resourceRegistrations, "Resource handlers list must not be null");
for (ResourceRegistration resource : resourceRegistrations) {
this.resources.put(resource.resource().name(), resource);
this.resources.put(resource.resource().uri(), resource);
}
return this;
}

29
mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java
View file

@ -153,31 +153,23 @@ public class StdioServerTransport implements McpTransport {
try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
if (!isClosing) {
logger.error("Failed to enqueue message");
}
logIfNotClosing("Failed to enqueue message");
break;
}
}
catch (Exception e) {
if (!isClosing) {
logger.error("Error processing inbound message", e);
}
logIfNotClosing("Error processing inbound message", e);
break;
}
}
catch (IOException e) {
if (!isClosing) {
logger.error("Error reading from stdin", e);
}
logIfNotClosing("Error reading from stdin", e);
break;
}
}
}
catch (Exception e) {
if (!isClosing) {
logger.error("Error in inbound processing", e);
}
logIfNotClosing("Error in inbound processing", e);
}
finally {
isClosing = true;
@ -254,7 +246,6 @@ public class StdioServerTransport implements McpTransport {
// Dispose schedulers with longer timeout
inboundScheduler.dispose();
outboundScheduler.dispose();
logger.info("Graceful shutdown completed");
}
catch (Exception e) {
@ -268,4 +259,16 @@ public class StdioServerTransport implements McpTransport {
return this.objectMapper.convertValue(data, typeRef);
}
private void logIfNotClosing(String message, Exception e) {
if (!this.isClosing) {
logger.error(message, e);
}
}
private void logIfNotClosing(String message) {
if (!this.isClosing) {
logger.error(message);
}
}
}

4
spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/client/ClientStdio.java
View file

@ -19,9 +19,9 @@ import org.springframework.ai.mcp.client.transport.ServerParameters;
import org.springframework.ai.mcp.client.transport.StdioClientTransport;
/**
* @author Christian Tzolov
* With stdio transport, the MCP server is automatically started by the client. But you
* have to build the server jar first: <pre>./mvnw clean install -DskipTests</pre>
*/
public class ClientStdio {
public static void main(String[] args) {

54
spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/server/McpServerConfig.java
View file

@ -27,18 +27,21 @@ import org.springframework.web.reactive.function.server.RouterFunction;
@Configuration
public class McpServerConfig {
// STDIO transport
@Bean
@ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "stdio")
public StdioServerTransport stdioServerTransport() {
return new StdioServerTransport();
}
// SSE transport
@Bean
@ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse")
public SseServerTransport sseServerTransport() {
return new SseServerTransport(new ObjectMapper(), "/mcp/message");
}
// Router function for SSE transport used by Spring WebFlux to start an HTTP server.
@Bean
@ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse")
public RouterFunction<?> mcpRouterFunction(SseServerTransport transport) {
@ -59,13 +62,13 @@ public class McpServerConfig {
return McpServer.using(transport)
.info("MCP Demo Server", "1.0.0")
.capabilities(capabilities)
.tools(List.of(weatherTool(), calculatorTool()))
.resources(Map.of("system://info", systemInfoResource()))
.prompts(List.of(greetingPrompt()))
.resources(systemInfoResourceRegistration())
.prompts(greetingPromptRegistration())
.tools(weatherToolRegistration(), calculatorToolRegistration())
.async();
} // @formatter:on
private static ResourceRegistration systemInfoResource() {
private static ResourceRegistration systemInfoResourceRegistration() {
// Create a resource registration for system information
var systemInfoResource = new McpSchema.Resource( // @formatter:off
@ -98,7 +101,26 @@ public class McpServerConfig {
return resourceRegistration;
}
private static ToolRegistration weatherTool() {
private static PromptRegistration greetingPromptRegistration() {
var prompt = new McpSchema.Prompt("greeting", "A friendly greeting prompt",
List.of(new McpSchema.PromptArgument("name", "The name to greet", true)));
return new PromptRegistration(prompt, getPromptRequest -> {
String nameArgument = (String) getPromptRequest.arguments().get("name");
if (nameArgument == null) {
nameArgument = "friend";
}
var userMessage = new PromptMessage(Role.USER,
new TextContent("Hello " + nameArgument + "! How can I assist you today?"));
return new GetPromptResult("A personalized greeting message", List.of(userMessage));
});
}
private static ToolRegistration weatherToolRegistration() {
return new ToolRegistration(
new McpSchema.Tool("weather", "Weather forecast tool by location", Map.of("city", "String")),
(arguments) -> {
@ -108,7 +130,7 @@ public class McpServerConfig {
});
}
private static ToolRegistration calculatorTool() {
private static ToolRegistration calculatorToolRegistration() {
return new ToolRegistration(new McpSchema.Tool("calculator",
"Performs basic arithmetic operations (add, subtract, multiply, divide)", """
{
@ -164,24 +186,4 @@ public class McpServerConfig {
});
}
private static PromptRegistration greetingPrompt() {
var nameArg = new McpSchema.PromptArgument("name", "The name to greet", true // required
// parameter
);
var prompt = new McpSchema.Prompt("greeting", "A friendly greeting prompt", List.of(nameArg));
return new PromptRegistration(prompt, request -> {
String name = (String) request.arguments().get("name");
if (name == null) {
name = "friend";
}
var assistantMessage = new PromptMessage(Role.ASSISTANT,
new TextContent("Hello " + name + "! How can I assist you today?"));
return new GetPromptResult("A personalized greeting message", List.of(assistantMessage));
});
}
}