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

fix: Improve MCP server stability and test reliability

Browse source 
- Replace bounded elastic schedulers with single thread executors
- Add timeouts to MCP server tests
- Enhance error handling in StdioServerTransport
- Comment out unused PaginatedRequest unmarshalling
- Simplify shutdown logic in StdioServerTransport
This commit is contained in:
Christian Tzolov 2024-12-30 12:15:39 +01:00
commit ad3cdd20a9
2 changed files with 50 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -349,9 +349,9 @@ public class McpAsyncServer {
private DefaultMcpSession.RequestHandler resourcesListRequestHandler() {
return params -> {
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.PaginatedRequest>() {
});
// McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
// new TypeReference<McpSchema.PaginatedRequest>() {
// });
var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList();
@ -361,9 +361,9 @@ public class McpAsyncServer {
private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() {
return params -> {
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.PaginatedRequest>() {
});
// McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
// new TypeReference<McpSchema.PaginatedRequest>() {
// });
return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
};

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

@ -23,6 +23,7 @@ import java.io.InputStreamReader;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.concurrent.Executors;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
@ -97,8 +98,8 @@ public class StdioServerTransport implements McpTransport {
this.outputStream = System.out;
// Use bounded schedulers for better resource management
this.inboundScheduler = Schedulers.newBoundedElastic(1, 1, "inbound");
this.outboundScheduler = Schedulers.newBoundedElastic(1, 1, "outbound");
this.inboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "inbound");
this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "outbound");
}
@Override
@ -139,38 +140,43 @@ public class StdioServerTransport implements McpTransport {
private void startInboundProcessing() {
this.inboundScheduler.schedule(() -> {
inboundReady.tryEmitValue(null);
try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
String line;
while (!isClosing && (line = reader.readLine()) != null) {
BufferedReader reader = null;
try {
reader = new BufferedReader(new InputStreamReader(inputStream));
while (!isClosing) {
try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
String line = reader.readLine();
if (line == null || isClosing) {
break;
}
try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
if (!isClosing) {
logger.error("Failed to enqueue message");
}
break;
}
}
catch (Exception e) {
if (!isClosing) {
logger.error("Failed to enqueue message");
logger.error("Error processing inbound message", e);
}
break;
}
}
catch (Exception e) {
catch (IOException e) {
if (!isClosing) {
logger.error("Error processing inbound message", e);
logger.error("Error reading from stdin", e);
}
break;
}
}
}
catch (IOException e) {
// Check isClosing before the error occurs to properly categorize it
boolean wasClosing = isClosing;
isClosing = true;
if (!wasClosing && e.getMessage().equals("Pipe closed")) {
logger.debug("Stream closed during shutdown", e);
}
else if (!wasClosing) {
logger.error("Error reading from stdin", e);
}
else {
logger.debug("Stream error during shutdown", e);
catch (Exception e) {
if (!isClosing) {
logger.error("Error in inbound processing", e);
}
}
finally {
@ -234,6 +240,7 @@ public class StdioServerTransport implements McpTransport {
@Override
public Mono<Void> closeGracefully() {
return Mono.fromRunnable(() -> {
isClosing = true;
logger.debug("Initiating graceful shutdown");
@ -244,18 +251,10 @@ public class StdioServerTransport implements McpTransport {
return Mono.delay(Duration.ofMillis(100));
})).then(Mono.fromRunnable(() -> {
try {
// Dispose schedulers first
// Dispose schedulers with longer timeout
inboundScheduler.dispose();
outboundScheduler.dispose();
// Wait for schedulers to terminate
if (!inboundScheduler.isDisposed()) {
inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
if (!outboundScheduler.isDisposed()) {
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
logger.info("Graceful shutdown completed");
}
catch (Exception e) {