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) Disable stdio server sink draining on close

Browse source
This commit is contained in:
Christian Tzolov 2024-12-30 09:23:42 +01:00
commit dfa7435ad8
1 changed files with 39 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -248,46 +248,50 @@ 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(() -> { })
// First complete the sinks to stop processing // .then(Mono.defer(() -> {
inboundSink.tryEmitComplete(); // // First complete the sinks to stop processing
outboundSink.tryEmitComplete(); // inboundSink.tryEmitComplete();
return Mono.delay(Duration.ofMillis(200)); // outboundSink.tryEmitComplete();
})).then(Mono.fromRunnable(() -> { // return Mono.delay(Duration.ofMillis(100));
try { // }))
// Dispose schedulers first .then(Mono.fromRunnable(() -> {
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));
}
// Only after schedulers are disposed, close the streams
try { try {
if (inputStream != System.in) { // Dispose schedulers first
inputStream.close(); inboundScheduler.dispose();
} outboundScheduler.dispose();
if (outputStream != System.out) {
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"); // Wait for schedulers to terminate
} if (!inboundScheduler.isDisposed()) {
catch (Exception e) { inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
logger.error("Error during graceful shutdown", e); }
} if (!outboundScheduler.isDisposed()) {
})).then().subscribeOn(Schedulers.boundedElastic()); outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
// Only after schedulers are disposed, close the streams
try {
if (inputStream != System.in) {
inputStream.close();
}
if (outputStream != System.out) {
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");
}
catch (Exception e) {
logger.error("Error during graceful shutdown", e);
}
}))
.then()
.subscribeOn(Schedulers.boundedElastic());
} }
@Override @Override