refactor(mcp): improve StdioServerTransport shutdown and error handling
- 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:
parent
2de1d637c0
commit
9b6a0b537b
1 changed files with 48 additions and 25 deletions
|
|
@ -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,16 +248,14 @@ 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(() -> {
|
||||||
|
|
||||||
.then(Mono.fromRunnable(() -> {
|
|
||||||
try {
|
try {
|
||||||
|
// Dispose schedulers first
|
||||||
inboundScheduler.dispose();
|
inboundScheduler.dispose();
|
||||||
outboundScheduler.dispose();
|
outboundScheduler.dispose();
|
||||||
|
|
||||||
|
|
@ -257,14 +267,27 @@ public class StdioServerTransport implements McpTransport {
|
||||||
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
|
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");
|
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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue