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(transport): Add cleanup logic for inbound processing in StdioServerTransport

Browse source 
Resolves #72

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov 2025-01-24 13:08:41 +01:00 • committed by Christian Tzolov
commit 48ab97bf23
1 changed files with 50 additions and 45 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -118,6 +118,11 @@ 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(() -> {
this.outboundSink.tryEmitComplete();
this.inboundScheduler.dispose();
this.outboundScheduler.dispose();
})
.subscribe(); .subscribe();
} }
@ -184,51 +189,51 @@ public class StdioServerTransport implements ServerMcpTransport {
*/ */
private void startOutboundProcessing() { private void startOutboundProcessing() {
Function<Flux<JSONRPCMessage>, Flux<JSONRPCMessage>> outboundConsumer = messages -> messages // @formatter:off Function<Flux<JSONRPCMessage>, Flux<JSONRPCMessage>> outboundConsumer = messages -> messages // @formatter:off
.doOnSubscribe(subscription -> outboundReady.tryEmitValue(null)) .doOnSubscribe(subscription -> outboundReady.tryEmitValue(null))
.publishOn(outboundScheduler) .publishOn(outboundScheduler)
.handle((message, sink) -> { .handle((message, sink) -> {
if (message != null && !isClosing) { if (message != null && !isClosing) {
try { try {
String jsonMessage = objectMapper.writeValueAsString(message); String jsonMessage = objectMapper.writeValueAsString(message);
// Escape any embedded newlines in the JSON message as per spec // Escape any embedded newlines in the JSON message as per spec
jsonMessage = jsonMessage.replace("\r\n", "\\n").replace("\n", "\\n").replace("\r", "\\n"); jsonMessage = jsonMessage.replace("\r\n", "\\n").replace("\n", "\\n").replace("\r", "\\n");
synchronized (outputStream) { synchronized (outputStream) {
outputStream.write(jsonMessage.getBytes(StandardCharsets.UTF_8)); outputStream.write(jsonMessage.getBytes(StandardCharsets.UTF_8));
outputStream.write("\n".getBytes(StandardCharsets.UTF_8)); outputStream.write("\n".getBytes(StandardCharsets.UTF_8));
outputStream.flush(); outputStream.flush();
} }
sink.next(message); sink.next(message);
} }
catch (IOException e) { catch (IOException e) {
if (!isClosing) { if (!isClosing) {
logger.error("Error writing message", e); logger.error("Error writing message", e);
sink.error(new RuntimeException(e)); sink.error(new RuntimeException(e));
} }
else { else {
logger.debug("Stream closed during shutdown", e); logger.debug("Stream closed during shutdown", e);
} }
} }
} }
else if (isClosing) { else if (isClosing) {
sink.complete(); sink.complete();
} }
}) })
.doOnComplete(() -> { .doOnComplete(() -> {
isClosing = true; isClosing = true;
outboundSink.tryEmitComplete(); outboundSink.tryEmitComplete();
}) })
.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(); outboundSink.tryEmitComplete();
} }
}) })
.map(msg -> (JSONRPCMessage) msg); .map(msg -> (JSONRPCMessage) msg);
outboundConsumer.apply(outboundSink.asFlux()).subscribe(); outboundConsumer.apply(outboundSink.asFlux()).subscribe();
} // @formatter:on } // @formatter:on
@Override @Override
public Mono<Void> closeGracefully() { public Mono<Void> closeGracefully() {