diff --git a/src/main/java/spring/ai/mcp/client/McpAsyncClient.java b/src/main/java/spring/ai/mcp/client/McpAsyncClient.java index 5defeba..0b2813d 100644 --- a/src/main/java/spring/ai/mcp/client/McpAsyncClient.java +++ b/src/main/java/spring/ai/mcp/client/McpAsyncClient.java @@ -1,22 +1,20 @@ package spring.ai.mcp.client; import java.time.Duration; -import java.util.concurrent.CompletableFuture; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import reactor.core.publisher.Mono; +import spring.ai.mcp.spec.McpAsyncTransport; import spring.ai.mcp.spec.McpSchema; -import spring.ai.mcp.spec.McpSession; -import spring.ai.mcp.spec.McpTransport; public class McpAsyncClient extends McpAsyncSession { - public McpAsyncClient(McpTransport transport) { + public McpAsyncClient(McpAsyncTransport transport) { this(transport, Duration.ofSeconds(10), new ObjectMapper()); } - public McpAsyncClient(McpTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { + public McpAsyncClient(McpAsyncTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { super(requestTimeout, objectMapper, transport); } diff --git a/src/main/java/spring/ai/mcp/client/McpAsyncSession.java b/src/main/java/spring/ai/mcp/client/McpAsyncSession.java index 7fde172..2ee02ad 100644 --- a/src/main/java/spring/ai/mcp/client/McpAsyncSession.java +++ b/src/main/java/spring/ai/mcp/client/McpAsyncSession.java @@ -3,18 +3,15 @@ package spring.ai.mcp.client; import java.time.Duration; import java.util.Map; import java.util.UUID; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.TimeUnit; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import reactor.core.publisher.Mono; import reactor.core.publisher.MonoSink; -import reactor.core.publisher.Sinks; -import spring.ai.mcp.spec.McpSchema; -import spring.ai.mcp.spec.McpSession; +import spring.ai.mcp.spec.McpAsyncTransport; import spring.ai.mcp.spec.McpTransport; +import spring.ai.mcp.spec.McpSchema; public class McpAsyncSession { @@ -25,11 +22,11 @@ public class McpAsyncSession { private final ObjectMapper objectMapper; - private final McpTransport transport; + private final McpAsyncTransport transport; public McpAsyncSession(Duration requestTimeout, ObjectMapper objectMapper, - McpTransport transport) { + McpAsyncTransport transport) { this.requestTimeout = requestTimeout; this.objectMapper = objectMapper; this.transport = transport; @@ -65,8 +62,15 @@ public class McpAsyncSession { this.pendingResponses.put(requestId, sink); McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, requestId, requestParams); try { - // TODO convert the transport to non-blocking as well - this.transport.sendMessage(jsonrpcRequest); + // TODO: This is non-blocking, but it's actually a synchronous call, + // perhaps there's no need to make it return Mono? + this.transport.sendMessage(jsonrpcRequest) + // TODO: It's most efficient to create a dedicated + // Subscriber here + .subscribe(v -> {}, e -> { + this.pendingResponses.remove(requestId); + sink.error(e); + }); } catch (Exception e) { sink.error(e); } diff --git a/src/main/java/spring/ai/mcp/client/McpClient.java b/src/main/java/spring/ai/mcp/client/McpClient.java index e59cee0..d85125a 100644 --- a/src/main/java/spring/ai/mcp/client/McpClient.java +++ b/src/main/java/spring/ai/mcp/client/McpClient.java @@ -20,6 +20,8 @@ import java.util.concurrent.CompletableFuture; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; +import spring.ai.mcp.spec.McpAsyncTransport; +import spring.ai.mcp.spec.McpTransport; import spring.ai.mcp.spec.McpSchema; import spring.ai.mcp.spec.McpSession; import spring.ai.mcp.spec.McpTransport; @@ -59,20 +61,20 @@ public class McpClient { // JSONtoPOJOCodec type // .sync(); - public static McpAsyncClient async(McpTransport transport) { + public static McpAsyncClient async(McpAsyncTransport transport) { return new McpAsyncClient(transport); } - public static McpAsyncClient async(McpTransport transport, Duration requestTimeout, + public static McpAsyncClient async(McpAsyncTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { return new McpAsyncClient(transport, requestTimeout, objectMapper); } - public static McpSyncClient sync(McpTransport transport) { + public static McpSyncClient sync(McpAsyncTransport transport) { return new McpSyncClient(async(transport)); } - public static McpSyncClient sync(McpTransport transport, Duration requestTimeout, + public static McpSyncClient sync(McpAsyncTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { return new McpSyncClient(async(transport, requestTimeout, objectMapper)); } diff --git a/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java b/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java index 2907fee..b4c275d 100644 --- a/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java +++ b/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java @@ -6,9 +6,10 @@ import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.List; -import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import reactor.core.scheduler.Scheduler; +import reactor.core.scheduler.Schedulers; import spring.ai.mcp.spec.DefaultMcpTransport; import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; import spring.ai.mcp.spec.McpSchema.JSONRPCNotification; @@ -31,7 +32,9 @@ public class StdioServerTransport extends DefaultMcpTransport { private BufferedWriter processWriter; - private ExecutorService executorService; + private Scheduler inboundScheduler; + private Scheduler outboundScheduler; + private Scheduler errorScheduler; private volatile boolean isRunning; @@ -49,7 +52,15 @@ public class StdioServerTransport extends DefaultMcpTransport { this.params = params; // Start threads - this.executorService = Executors.newFixedThreadPool(3); + this.inboundScheduler = + Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), + "inbound"); + this.outboundScheduler = + Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), + "outbound"); + this.errorScheduler = + Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), + "error"); } @Override @@ -83,9 +94,9 @@ public class StdioServerTransport extends DefaultMcpTransport { // Start threads this.isRunning = true; - startReaderThread(); - startWriterThread(); - startErrorReaderThread(); + startInboundProcessing(); + startOutboundProcessing(); + startErrorProcessing(); } @@ -98,9 +109,8 @@ public class StdioServerTransport extends DefaultMcpTransport { } } - private void startErrorReaderThread() { - - this.executorService.execute(() -> { + private void startErrorProcessing() { + this.errorScheduler.schedule(() -> { try { String line; while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) { @@ -143,15 +153,15 @@ public class StdioServerTransport extends DefaultMcpTransport { throw new IllegalArgumentException("Cannot deserialize JSONRPCMessage: " + jsonText); } - private void startReaderThread() { - this.executorService.execute(() -> { + private void startInboundProcessing() { + this.inboundScheduler.schedule(() -> { try { String line; while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) { try { JSONRPCMessage message = deserializeJsonRpcMessage(line); if (!this.getInboundSink().tryEmitNext(message).isSuccess()) { - // TODO: Back off, reschedule -> same comment as for writing + // TODO: Back off, reschedule, give up? throw new RuntimeException("Failed to enqueue message"); } } @@ -171,29 +181,33 @@ public class StdioServerTransport extends DefaultMcpTransport { }); } - private void startWriterThread() { - this.executorService.execute(() -> { - this.getOutboundSink().asFlux().handle((message, s) -> { - if (message != null) { - try { - this.processWriter.write(ModelOptionsUtils.toJsonString(message)); - this.processWriter.newLine(); - this.processWriter.flush(); - s.next(message); - } catch (IOException e) { - s.error(new RuntimeException(e)); - } - } - }).subscribe(); - }); + private void startOutboundProcessing() { + this.getOutboundSink() + .asFlux() + // this bit is important since writes come from user threads and we + // want to ensure that the actual writing happens on a dedicated thread + .publishOn(outboundScheduler) + .handle((message, s) -> { + if (message != null) { + try { + this.processWriter.write(ModelOptionsUtils.toJsonString(message)); + this.processWriter.newLine(); + this.processWriter.flush(); + s.next(message); + } + catch (IOException e) { + s.error(new RuntimeException(e)); + } + } + }) + .subscribe(); } + // TODO: provide a non-blocking variant with graceful option public void stop() { - - this.isRunning = false; - - // Interrupt threads - this.executorService.shutdownNow(); + this.outboundScheduler.dispose(); + this.inboundScheduler.dispose(); + this.errorScheduler.dispose(); // Destroy process if (this.process != null) { @@ -201,10 +215,6 @@ public class StdioServerTransport extends DefaultMcpTransport { } } - public boolean isRunning() { - return this.isRunning; - } - @Override public void close() { stop(); diff --git a/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java b/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java index 3105368..62e476b 100644 --- a/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java +++ b/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java @@ -16,10 +16,9 @@ package spring.ai.mcp.spec; import java.time.Duration; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.function.Consumer; +import reactor.core.publisher.Mono; import reactor.core.publisher.Sinks; import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; @@ -27,7 +26,7 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; * @author Christian Tzolov * @since 1.0.0 */ -public class DefaultMcpTransport implements McpTransport { +public class DefaultMcpTransport implements McpAsyncTransport { private final Sinks.Many errorSink; @@ -35,8 +34,6 @@ public class DefaultMcpTransport implements McpTransport { private final Sinks.Many outboundSink; - private final ExecutorService executorService; - private final Duration writeTimeout; private Consumer messageReader = message -> System.out.println("Received message: " + message); @@ -56,23 +53,18 @@ public class DefaultMcpTransport implements McpTransport { this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.writeTimeout = readTimeout; - this.executorService = Executors.newFixedThreadPool(2); this.handleIncomingMessages(); } private void handleIncomingMessages() { - this.executorService.execute(() -> { - this.inboundSink.asFlux() - .subscribe(message -> this.messageReader.accept(message)); - }); + this.inboundSink.asFlux() + .subscribe(message -> this.messageReader.accept(message)); } private void handleIncomingErrors() { - this.executorService.execute(() -> { - this.errorSink.asFlux().subscribe(e -> { - this.errorReader.accept(e); - System.err.println(e); - }); + this.errorSink.asFlux().subscribe(e -> { + this.errorReader.accept(e); + System.err.println(e); }); } @@ -106,15 +98,19 @@ public class DefaultMcpTransport implements McpTransport { @Override public void close() { // this.onClose(); - this.executorService.shutdownNow(); } - public void sendMessage(JSONRPCMessage message) { - if (!this.outboundSink.tryEmitNext(message).isSuccess()) { + @Override + public Mono sendMessage(JSONRPCMessage message) { + 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"); + // In this approach we delegate the retry and the backpressure onto the + // caller. This might be enough for most cases. + return Mono.empty(); + } else { + return Mono.error(new RuntimeException("Failed to enqueue message")); } } } diff --git a/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java b/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java new file mode 100644 index 0000000..b8109c7 --- /dev/null +++ b/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java @@ -0,0 +1,39 @@ +/* +* Copyright 2024 - 2024 the original author or authors. +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* https://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ +package spring.ai.mcp.spec; + +import java.util.function.Consumer; + +import reactor.core.publisher.Mono; +import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; + +/** + * @author Christian Tzolov + * @since 1.0.0 + */ + +public interface McpAsyncTransport extends AutoCloseable{ + + void start(); + + void close(); + + void setMessageHandler(Consumer messageReader); + + void setErrorHandler(Consumer errorReader); + + Mono sendMessage(JSONRPCMessage message); +}