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.List;
import java.util.concurrent.ExecutorService; import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors; import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import spring.ai.mcp.spec.DefaultMcpTransport; import spring.ai.mcp.spec.DefaultMcpTransport;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
@ -107,7 +106,8 @@ public class StdioServerTransport extends DefaultMcpTransport {
while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) { while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) {
try { try {
System.out.println("Received error line: " + line); 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) { catch (Exception e) {
throw new RuntimeException(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) { while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) {
try { try {
JSONRPCMessage message = deserializeJsonRpcMessage(line); 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) { catch (Exception e) {
throw new RuntimeException(e); throw new RuntimeException(e);
@ -170,28 +173,18 @@ public class StdioServerTransport extends DefaultMcpTransport {
private void startWriterThread() { private void startWriterThread() {
this.executorService.execute(() -> { this.executorService.execute(() -> {
try { this.getOutboundSink().asFlux().handle((message, s) -> {
while (this.isRunning && this.processWriter != null && !Thread.currentThread().isInterrupted()) { if (message != null) {
try { try {
JSONRPCMessage message = this.getDataWriteQueue().poll(100, TimeUnit.MILLISECONDS); this.processWriter.write(ModelOptionsUtils.toJsonString(message));
if (message != null) { this.processWriter.newLine();
this.processWriter.write(ModelOptionsUtils.toJsonString(message)); this.processWriter.flush();
this.processWriter.newLine(); s.next(message);
this.processWriter.flush(); } catch (IOException e) {
} s.error(new RuntimeException(e));
}
catch (InterruptedException e) {
Thread.currentThread().interrupt();
break;
} }
} }
} }).subscribe();
catch (IOException e) {
throw new RuntimeException(e);
}
finally {
this.isRunning = false;
}
}); });
} }

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

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