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(mcp): Improve transport shutdown and error handling

Browse source 
- Add isClosing flag to coordinate graceful shutdown
- Synchronize output stream writes to prevent race conditions
- Enhance error handling in stream processing
- Add delay for pending message processing during shutdown
- Improve connection state tracking in tests
This commit is contained in:
Christian Tzolov 2024-12-29 10:24:32 +01:00
commit 9130e742fd
3 changed files with 124 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

101
mcp/src/main/java/org/springframework/ai/mcp/client/transport/StdioClientTransport.java
View file

@ -20,6 +20,7 @@ import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
@ -77,6 +78,8 @@ public class StdioClientTransport implements McpTransport {
private final Sinks.Many<String> errorSink;
private volatile boolean isClosing = false;
// visible for tests
private Consumer<String> errorHandler = error -> logger.error("Error received: {}", error);
@ -200,19 +203,32 @@ public class StdioClientTransport implements McpTransport {
try (BufferedReader processErrorReader = new BufferedReader(
new InputStreamReader(process.getErrorStream()))) {
String line;
while ((line = processErrorReader.readLine()) != null) {
while (!isClosing && (line = processErrorReader.readLine()) != null) {
try {
logger.error("Received error line: {}", line);
// TODO: handle errors, etc.
this.errorSink.tryEmitNext(line);
if (!this.errorSink.tryEmitNext(line).isSuccess()) {
if (!isClosing) {
logger.error("Failed to emit error message");
}
break;
}
}
catch (Exception e) {
throw new RuntimeException(e);
if (!isClosing) {
logger.error("Error processing error message", e);
}
break;
}
}
}
catch (IOException e) {
throw new RuntimeException(e);
if (!isClosing) {
logger.error("Error reading from error stream", e);
}
}
finally {
isClosing = true;
errorSink.tryEmitComplete();
}
});
}
@ -254,21 +270,32 @@ public class StdioClientTransport implements McpTransport {
this.inboundScheduler.schedule(() -> {
try (BufferedReader processReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
String line;
while ((line = processReader.readLine()) != null) {
while (!isClosing && (line = processReader.readLine()) != null) {
try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
// TODO: Back off, reschedule, give up?
throw new RuntimeException("Failed to enqueue message");
if (!isClosing) {
logger.error("Failed to enqueue inbound message");
}
break;
}
}
catch (Exception e) {
throw new RuntimeException(e);
if (!isClosing) {
logger.error("Error processing inbound message", e);
}
break;
}
}
}
catch (IOException e) {
throw new RuntimeException(e);
if (!isClosing) {
logger.error("Error reading from input stream", e);
}
}
finally {
isClosing = true;
inboundSink.tryEmitComplete();
}
});
}
@ -284,7 +311,7 @@ public class StdioClientTransport implements McpTransport {
// want to ensure that the actual writing happens on a dedicated thread
.publishOn(outboundScheduler)
.handle((message, s) -> {
if (message != null) {
if (message != null && !isClosing) {
try {
String jsonMessage = objectMapper.writeValueAsString(message);
// Escape any embedded newlines in the JSON message as per spec:
@ -293,9 +320,12 @@ public class StdioClientTransport implements McpTransport {
// embedded newlines.
jsonMessage = jsonMessage.replace("\r\n", "\\n").replace("\n", "\\n").replace("\r", "\\n");
this.process.getOutputStream().write(jsonMessage.getBytes(StandardCharsets.UTF_8));
this.process.getOutputStream().write("\n".getBytes(StandardCharsets.UTF_8));
this.process.getOutputStream().flush();
var os = this.process.getOutputStream();
synchronized (os) {
os.write(jsonMessage.getBytes(StandardCharsets.UTF_8));
os.write("\n".getBytes(StandardCharsets.UTF_8));
os.flush();
}
s.next(message);
}
catch (IOException e) {
@ -306,7 +336,16 @@ public class StdioClientTransport implements McpTransport {
}
protected void handleOutbound(Function<Flux<JSONRPCMessage>, Flux<JSONRPCMessage>> outboundConsumer) {
outboundConsumer.apply(outboundSink.asFlux()).subscribe();
outboundConsumer.apply(outboundSink.asFlux()).doOnComplete(() -> {
isClosing = true;
outboundSink.tryEmitComplete();
}).doOnError(e -> {
if (!isClosing) {
logger.error("Error in outbound processing", e);
isClosing = true;
outboundSink.tryEmitComplete();
}
}).subscribe();
}
/**
@ -317,7 +356,18 @@ public class StdioClientTransport implements McpTransport {
*/
@Override
public Mono<Void> closeGracefully() {
return Mono.fromFuture(() -> {
return Mono.fromRunnable(() -> {
isClosing = true;
logger.debug("Initiating graceful shutdown");
}).then(Mono.defer(() -> {
// First complete all sinks to stop accepting new messages
inboundSink.tryEmitComplete();
outboundSink.tryEmitComplete();
errorSink.tryEmitComplete();
// Give a short time for any pending messages to be processed
return Mono.delay(Duration.ofMillis(100));
})).then(Mono.fromFuture(() -> {
logger.info("Sending TERM to process");
if (this.process != null) {
this.process.destroy();
@ -326,16 +376,23 @@ public class StdioClientTransport implements McpTransport {
else {
return CompletableFuture.failedFuture(new RuntimeException("Process not started"));
}
}).doOnNext(process -> {
})).doOnNext(process -> {
if (process.exitValue() != 0) {
logger.warn("Process terminated with code " + process.exitValue());
}
}).then(Mono.fromRunnable(() -> {
// The Threads are blocked on readLine so disposeGracefully would not
// interrupt them, therefore we issue an async hard dispose.
inboundScheduler.dispose();
errorScheduler.dispose();
outboundScheduler.dispose();
try {
// The Threads are blocked on readLine so disposeGracefully would not
// interrupt them, therefore we issue an async hard dispose.
inboundScheduler.dispose();
errorScheduler.dispose();
outboundScheduler.dispose();
logger.info("Graceful shutdown completed");
}
catch (Exception e) {
logger.error("Error during graceful shutdown", e);
}
})).then().subscribeOn(Schedulers.boundedElastic());
}

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

@ -176,9 +176,7 @@ public class StdioServerTransport implements McpTransport {
}
}
finally {
if (!isClosing) {
isClosing = true;
}
isClosing = true;
inboundSink.tryEmitComplete();
}
});
@ -218,9 +216,7 @@ public class StdioServerTransport implements McpTransport {
}
})
.doOnComplete(() -> {
if (!isClosing) {
isClosing = true;
}
isClosing = true;
outboundSink.tryEmitComplete();
})
.doOnError(e -> {
@ -240,29 +236,35 @@ public class StdioServerTransport implements McpTransport {
return Mono.fromRunnable(() -> {
isClosing = true;
logger.debug("Initiating graceful shutdown");
// }).then(Mono.delay(Duration.ofMillis(100))).then(Mono.fromRunnable(() -> {
}).then(Mono.fromRunnable(() -> {
try {
// inboundSink.tryEmitComplete();
// outboundSink.tryEmitComplete();
})
// .then(Mono.defer(() -> {
// inboundSink.tryEmitComplete();
// outboundSink.tryEmitComplete();
// return Mono.delay(Duration.ofMillis(100));
// }))
inboundScheduler.dispose();
outboundScheduler.dispose();
.then(Mono.fromRunnable(() -> {
try {
// Wait for schedulers to terminate
if (!inboundScheduler.isDisposed()) {
inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
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");
}
if (!outboundScheduler.isDisposed()) {
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
catch (Exception e) {
logger.error("Error during graceful shutdown", e);
}
logger.info("Graceful shutdown completed");
}
catch (Exception e) {
logger.error("Error during graceful shutdown", e);
}
})).then().subscribeOn(Schedulers.boundedElastic());
}))
.then()
.subscribeOn(Schedulers.boundedElastic());
}
@Override

20
mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java
View file

@ -49,11 +49,11 @@ class McpAsyncClientResponseHandlerTests {
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
private Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
private final Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
private Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
private final Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
private Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
if (inbound.tryEmitNext(message).isFailure()) {
@ -82,17 +82,29 @@ class McpAsyncClientResponseHandlerTests {
return outboundView.blockFirst();
}
private volatile boolean connected = false;
@Override
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
if (connected) {
return Mono.error(new IllegalStateException("Already connected"));
}
connected = true;
return inbound.asFlux()
.publishOn(Schedulers.boundedElastic())
.flatMap(message -> Mono.just(message).transform(handler))
.doFinally(signal -> connected = false)
.then();
}
@Override
public Mono<Void> closeGracefully() {
return Mono.empty();
return Mono.defer(() -> {
connected = false;
outgoing.tryEmitComplete();
inbound.tryEmitComplete();
return Mono.empty();
});
}
@Override