non-blocking sinks processing on dedicated Schedulers

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);
 	}
 

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.McpAsyncTransport;
-import spring.ai.mcp.spec.McpSchema;
-import spring.ai.mcp.spec.McpSession;
 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
+				// TODO: This is non-blocking, but it's actually a synchronous call,
-				this.transport.sendMessage(jsonrpcRequest);
+				// 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);
 			}

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));
 	}

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();
+		startInboundProcessing();
-		startWriterThread();
+		startOutboundProcessing();
-		startErrorReaderThread();
+		startErrorProcessing();
 
 	}
 
@@ -98,9 +109,8 @@ public class StdioServerTransport extends DefaultMcpTransport {
 		}
 	}
 
-	private void startErrorReaderThread() {
+	private void startErrorProcessing() {
-
+		this.errorScheduler.schedule(() -> {
-		this.executorService.execute(() -> {
 			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() {
+	private void startInboundProcessing() {
-		this.executorService.execute(() -> {
+		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() {
+	private void startOutboundProcessing() {
-		this.executorService.execute(() -> {
+		this.getOutboundSink()
-			this.getOutboundSink().asFlux().handle((message, s) -> {
+		    .asFlux()
-				if (message != null) {
+		    // this bit is important since writes come from user threads and we
-					try {
+		    // want to ensure that the actual writing happens on a dedicated thread
-						this.processWriter.write(ModelOptionsUtils.toJsonString(message));
+		    .publishOn(outboundScheduler)
-						this.processWriter.newLine();
+		    .handle((message, s) -> {
-						this.processWriter.flush();
+			    if (message != null) {
-						s.next(message);
+				    try {
-					} catch (IOException e) {
+					    this.processWriter.write(ModelOptionsUtils.toJsonString(message));
-						s.error(new RuntimeException(e));
+					    this.processWriter.newLine();
-					}
+					    this.processWriter.flush();
-				}
+					    s.next(message);
-			}).subscribe();
+				    }
-		});
+				    catch (IOException e) {
+					    s.error(new RuntimeException(e));
+				    }
+			    }
+		    })
+		    .subscribe();
 	}
 
+	// TODO: provide a non-blocking variant with graceful option
 	public void stop() {
-		
+		this.outboundScheduler.dispose();
-		this.isRunning = false;
+		this.inboundScheduler.dispose();
-
+		this.errorScheduler.dispose();
-		// Interrupt threads
-		this.executorService.shutdownNow();
 
 		// 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();

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<String> errorSink;
 
@@ -35,8 +34,6 @@ public class DefaultMcpTransport implements McpTransport {
 
 	private final Sinks.Many<JSONRPCMessage> outboundSink;
 
-	private final ExecutorService executorService;
-
 	private final Duration writeTimeout;
 
 	private Consumer<JSONRPCMessage> 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()
-			this.inboundSink.asFlux()
+		                .subscribe(message -> this.messageReader.accept(message));
-			                .subscribe(message -> this.messageReader.accept(message));
-		});
 	}
 
 	private void handleIncomingErrors() {
-		this.executorService.execute(() -> {
+		this.errorSink.asFlux().subscribe(e -> {
-			this.errorSink.asFlux().subscribe(e -> {
+			this.errorReader.accept(e);
-				this.errorReader.accept(e);
+			System.err.println(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) {
+	@Override
-		if (!this.outboundSink.tryEmitNext(message).isSuccess()) {
+	public Mono<Void> 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"));
 		}
 	}
 }

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<JSONRPCMessage> messageReader);
+
+	void setErrorHandler(Consumer<String> errorReader);
+
+	Mono<Void> sendMessage(JSONRPCMessage message);    
+}