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

This commit is messy but it reworks the transport to work with Sinks and be non-blocking on the caller side

Browse source
This commit is contained in:
Dariusz Jędrzejczyk 2024-12-06 23:31:13 +01:00
commit d21565a559
No known key found for this signature in database
GPG key ID: FD786A4E2335C448
2 changed files with 45 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

37
src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java
View file

@ -8,7 +8,6 @@ import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import spring.ai.mcp.spec.DefaultMcpTransport;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
@ -107,7 +106,8 @@ public class StdioServerTransport extends DefaultMcpTransport {
while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) {
try {
System.out.println("Received error line: " + line);
this.getErrorReadQueue().offer(line, 100, TimeUnit.MILLISECONDS);
// TODO: handle errors, etc.
this.getErrorSink().tryEmitNext(line);
}
catch (Exception e) {
throw new RuntimeException(e);
@ -150,7 +150,10 @@ public class StdioServerTransport extends DefaultMcpTransport {
while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) {
try {
JSONRPCMessage message = deserializeJsonRpcMessage(line);
this.getDataReadQueue().offer(message, 100, TimeUnit.MILLISECONDS);
if (!this.getInboundSink().tryEmitNext(message).isSuccess()) {
// TODO: Back off, reschedule -> same comment as for writing
throw new RuntimeException("Failed to enqueue message");
}
}
catch (Exception e) {
throw new RuntimeException(e);
@ -170,28 +173,18 @@ public class StdioServerTransport extends DefaultMcpTransport {
private void startWriterThread() {
this.executorService.execute(() -> {
try {
while (this.isRunning && this.processWriter != null && !Thread.currentThread().isInterrupted()) {
this.getOutboundSink().asFlux().handle((message, s) -> {
if (message != null) {
try {
JSONRPCMessage message = this.getDataWriteQueue().poll(100, TimeUnit.MILLISECONDS);
if (message != null) {
this.processWriter.write(ModelOptionsUtils.toJsonString(message));
this.processWriter.newLine();
this.processWriter.flush();
}
}
catch (InterruptedException e) {
Thread.currentThread().interrupt();
break;
this.processWriter.write(ModelOptionsUtils.toJsonString(message));
this.processWriter.newLine();
this.processWriter.flush();
s.next(message);
} catch (IOException e) {
s.error(new RuntimeException(e));
}
}
}
catch (IOException e) {
throw new RuntimeException(e);
}
finally {
this.isRunning = false;
}
}).subscribe();
});
}

71
src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java
View file

@ -16,13 +16,11 @@
package spring.ai.mcp.spec;
import java.time.Duration;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import reactor.core.publisher.Sinks;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
/**
@ -31,11 +29,11 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
*/
public class DefaultMcpTransport implements McpTransport {
private final BlockingQueue<String> errorReadQueue;
private final Sinks.Many<String> errorSink;
private final BlockingQueue<JSONRPCMessage> dataReadQueue;
private final Sinks.Many<JSONRPCMessage> inboundSink;
private final BlockingQueue<JSONRPCMessage> dataWriteQueue;
private final Sinks.Many<JSONRPCMessage> outboundSink;
private final ExecutorService executorService;
@ -47,9 +45,15 @@ public class DefaultMcpTransport implements McpTransport {
public DefaultMcpTransport(Duration readTimeout) {
this.errorReadQueue = new LinkedBlockingQueue<>();
this.dataReadQueue = new LinkedBlockingQueue<>();
this.dataWriteQueue = new LinkedBlockingQueue<>();
// TODO: consider the effects of buffering here -> the inter-process pipes are
// independent and the notifications can flood the client/server.
// Potentially, the interest in reading could be communicated from one party
// to the other so the Blocking IO Threads can pause consuming the stream
// buffers when there is no expectation for reading.
this.errorSink = Sinks.many().unicast().onBackpressureBuffer();
this.inboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.outboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.writeTimeout = readTimeout;
this.executorService = Executors.newFixedThreadPool(2);
@ -58,43 +62,30 @@ public class DefaultMcpTransport implements McpTransport {
private void handleIncomingMessages() {
this.executorService.execute(() -> {
while (!Thread.currentThread().isInterrupted()) {
try {
JSONRPCMessage message = this.dataReadQueue.take();
this.messageReader.accept(message);
}
catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
this.inboundSink.asFlux()
.subscribe(message -> this.messageReader.accept(message));
});
}
private void handleIncomingErrors() {
this.executorService.execute(() -> {
while (!Thread.currentThread().isInterrupted()) {
try {
String error = this.errorReadQueue.take();
this.errorReader.accept(error);
System.err.println(error);
}
catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
this.errorSink.asFlux().subscribe(e -> {
this.errorReader.accept(e);
System.err.println(e);
});
});
}
protected BlockingQueue<JSONRPCMessage> getDataReadQueue() {
return dataReadQueue;
protected Sinks.Many<JSONRPCMessage> getInboundSink() {
return inboundSink;
}
protected BlockingQueue<JSONRPCMessage> getDataWriteQueue() {
return dataWriteQueue;
protected Sinks.Many<JSONRPCMessage> getOutboundSink() {
return outboundSink;
}
protected BlockingQueue<String> getErrorReadQueue() {
return errorReadQueue;
protected Sinks.Many<String> getErrorSink() {
return errorSink;
}
public void setMessageHandler(Consumer<JSONRPCMessage> messageReader) {
@ -119,13 +110,11 @@ public class DefaultMcpTransport implements McpTransport {
}
public void sendMessage(JSONRPCMessage message) {
// Use offer with timeout to prevent blocking indefinitely
try {
if (!this.dataWriteQueue.offer(message, writeTimeout.toMillis(), TimeUnit.MILLISECONDS)) {
throw new RuntimeException("Failed to enqueue message");
}
} catch (InterruptedException e) {
throw new RuntimeException(e);
if (!this.outboundSink.tryEmitNext(message).isSuccess()) {
// TODO: essentially we could reschedule ourselves in some time and make
// another attempt with the already read data but pause reading until
// success
throw new RuntimeException("Failed to enqueue message");
}
}
}