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

refactor(mcp): improve StdioServerTransport shutdown and error handling

Browse source 
- Add specific error categorization for pipe closed errors
- Enhance shutdown sequence with proper stream and scheduler cleanup
- Add debug logging for shutdown-related stream errors
- Improve isClosing state management
This commit is contained in:
Christian Tzolov 2024-12-29 21:40:51 +01:00
commit 9b6a0b537b
1 changed files with 48 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -171,9 +171,18 @@ public class StdioServerTransport implements McpTransport {
} }
} }
catch (IOException e) { catch (IOException e) {
if (!isClosing) { // 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); logger.error("Error reading from stdin", e);
} }
else {
logger.debug("Stream error during shutdown", e);
}
} }
finally { finally {
isClosing = true; isClosing = true;
@ -209,6 +218,9 @@ public class StdioServerTransport implements McpTransport {
logger.error("Error writing message", e); logger.error("Error writing message", e);
sink.error(new RuntimeException(e)); sink.error(new RuntimeException(e));
} }
else {
logger.debug("Stream closed during shutdown", e);
}
} }
} }
else if (isClosing) { else if (isClosing) {
@ -236,35 +248,46 @@ public class StdioServerTransport implements McpTransport {
return Mono.fromRunnable(() -> { return Mono.fromRunnable(() -> {
isClosing = true; isClosing = true;
logger.debug("Initiating graceful shutdown"); logger.debug("Initiating graceful shutdown");
}) }).then(Mono.defer(() -> {
// .then(Mono.defer(() -> { // First complete the sinks to stop processing
// inboundSink.tryEmitComplete(); inboundSink.tryEmitComplete();
// outboundSink.tryEmitComplete(); outboundSink.tryEmitComplete();
// return Mono.delay(Duration.ofMillis(100)); return Mono.delay(Duration.ofMillis(100));
// })) })).then(Mono.fromRunnable(() -> {
try {
// Dispose schedulers first
inboundScheduler.dispose();
outboundScheduler.dispose();
.then(Mono.fromRunnable(() -> { // Wait for schedulers to terminate
if (!inboundScheduler.isDisposed()) {
inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
if (!outboundScheduler.isDisposed()) {
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
// Only after schedulers are disposed, close the streams
try { try {
if (inputStream != System.in) {
inboundScheduler.dispose(); inputStream.close();
outboundScheduler.dispose();
// Wait for schedulers to terminate
if (!inboundScheduler.isDisposed()) {
inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
} }
if (!outboundScheduler.isDisposed()) { if (outputStream != System.out) {
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5)); outputStream.flush();
outputStream.close();
} }
}
catch (IOException e) {
// Log but don't throw since we're shutting down
logger.debug("Error closing streams during shutdown", e);
}
logger.info("Graceful shutdown completed"); logger.info("Graceful shutdown completed");
} }
catch (Exception e) { catch (Exception e) {
logger.error("Error during graceful shutdown", e); logger.error("Error during graceful shutdown", e);
} }
})) })).then().subscribeOn(Schedulers.boundedElastic());
.then()
.subscribeOn(Schedulers.boundedElastic());
} }
@Override @Override