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

Reorder StdioServerTransport shutdown procedure to avoid errors

Browse source 
Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Dariusz Jędrzejczyk 2025-01-28 15:53:56 +01:00 • committed by Christian Tzolov
commit c8766098f6
1 changed files with 10 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

31
mcp/src/main/java/org/modelcontextprotocol/server/transport/StdioServerTransport.java
View file

@ -105,10 +105,10 @@ public class StdioServerTransport implements ServerMcpTransport {
.flatMap(message -> Mono.just(message) .flatMap(message -> Mono.just(message)
.transform(inboundMessageHandler) .transform(inboundMessageHandler)
.contextWrite(ctx -> ctx.put("observation", "myObservation"))) .contextWrite(ctx -> ctx.put("observation", "myObservation")))
.doOnComplete(() -> { .doOnTerminate(() -> {
// The outbound processing will dispose its scheduler upon completion
this.outboundSink.tryEmitComplete(); this.outboundSink.tryEmitComplete();
this.inboundScheduler.dispose(); this.inboundScheduler.dispose();
this.outboundScheduler.dispose();
}) })
.subscribe(); .subscribe();
} }
@ -208,13 +208,13 @@ public class StdioServerTransport implements ServerMcpTransport {
}) })
.doOnComplete(() -> { .doOnComplete(() -> {
isClosing = true; isClosing = true;
outboundSink.tryEmitComplete(); outboundScheduler.dispose();
}) })
.doOnError(e -> { .doOnError(e -> {
if (!isClosing) { if (!isClosing) {
logger.error("Error in outbound processing", e); logger.error("Error in outbound processing", e);
isClosing = true; isClosing = true;
outboundSink.tryEmitComplete(); outboundScheduler.dispose();
} }
}) })
.map(msg -> (JSONRPCMessage) msg); .map(msg -> (JSONRPCMessage) msg);
@ -224,26 +224,15 @@ public class StdioServerTransport implements ServerMcpTransport {
@Override @Override
public Mono<Void> closeGracefully() { public Mono<Void> closeGracefully() {
return Mono.<Void>defer(() -> {
return Mono.fromRunnable(() -> {
isClosing = true; isClosing = true;
logger.debug("Initiating graceful shutdown"); logger.debug("Initiating graceful shutdown");
}).then(Mono.defer(() -> { // Completing the inbound causes the outbound to be completed as well, so
// First complete the sinks to stop processing // we only close the inbound.
inboundSink.tryEmitComplete(); inboundSink.tryEmitComplete();
outboundSink.tryEmitComplete(); logger.info("Graceful shutdown complete");
return Mono.delay(Duration.ofMillis(100)); return Mono.empty();
})).then(Mono.fromRunnable(() -> { }).subscribeOn(Schedulers.boundedElastic());
try {
// Dispose schedulers with longer timeout
inboundScheduler.dispose();
outboundScheduler.dispose();
logger.info("Graceful shutdown completed");
}
catch (Exception e) {
logger.error("Error during graceful shutdown", e);
}
})).then().subscribeOn(Schedulers.boundedElastic());
} }
@Override @Override