refactor(transport): improve MCP transport layer design

- Rename DefaultMcpTransport to AbstractMcpTransport and make it abstract
- Add graceful shutdown support via closeGracefully() method
- Enhance McpTransport interface with default close() implementation
- Improve documentation across transport layer classes

mcp/pom.xml

@@ -24,7 +24,6 @@
 	<properties>
 		<java.version>21</java.version>
 		<maven.compiler.release>21</maven.compiler.release>
-		<!-- <spring-ai.version>1.0.0-M4</spring-ai.version> -->
 	</properties>
 	<dependencies>
 		<dependency>

mcp/src/main/java/spring/ai/experimental/mcp/client/McpClient.java

@@ -21,14 +21,30 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import spring.ai.experimental.mcp.spec.McpTransport;
 
 /**
- * The MCP client is the main entry point for interacting with the Model Context Protocol
+ * Factory class providing static methods for creating Model Context Protocol (MCP) clients.
- * (MCP) server.
+ * This class serves as the main entry point for establishing connections with MCP servers,
+ * offering both synchronous and asynchronous client implementations.
+ * 
+ * <p>The class provides factory methods to create either:
+ * <ul>
+ *   <li>{@link McpAsyncClient} for non-blocking operations
+ *   <li>{@link McpSyncClient} for blocking operations
+ * </ul>
+ * 
+ * <p>Each client type can be instantiated with default settings or with custom configuration
+ * including request timeout and JSON object mapping.
+ * 
+ * <p>Future implementations will introduce a builder pattern for more flexible client configuration.
  * 
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
 public class McpClient {
 
+    /**
+     * Private constructor to prevent instantiation as this is a utility class
+     * containing only static factory methods.
+     */
     private McpClient() {
     }
 
@@ -39,23 +55,51 @@ public class McpClient {
     //                                               JSONtoPOJOCodec type
     //           .sync();
 
+    /**
+     * Creates an asynchronous MCP client with default configuration.
+     *
+     * @param transport The transport layer implementation for MCP communication
+     * @return A new instance of {@link McpAsyncClient}
+     */
     public static McpAsyncClient async(McpTransport transport) {
         return new McpAsyncClient(transport);
     }
 
+    /**
+     * Creates an asynchronous MCP client with custom configuration.
+     *
+     * @param transport The transport layer implementation for MCP communication
+     * @param requestTimeout The duration to wait before timing out requests
+     * @param objectMapper The Jackson object mapper for JSON serialization/deserialization
+     * @return A new instance of {@link McpAsyncClient}
+     */
     public static McpAsyncClient async(McpTransport transport, Duration requestTimeout,
             ObjectMapper objectMapper) {
         return new McpAsyncClient(transport, requestTimeout, objectMapper);
     }
 
+    /**
+     * Creates a synchronous MCP client with default configuration.
+     * This method wraps an asynchronous client to provide synchronous operations.
+     *
+     * @param transport The transport layer implementation for MCP communication
+     * @return A new instance of {@link McpSyncClient}
+     */
     public static McpSyncClient sync(McpTransport transport) {
         return new McpSyncClient(async(transport));
     }
 
+    /**
+     * Creates a synchronous MCP client with custom configuration.
+     * This method wraps an asynchronous client to provide synchronous operations.
+     *
+     * @param transport The transport layer implementation for MCP communication
+     * @param requestTimeout The duration to wait before timing out requests
+     * @param objectMapper The Jackson object mapper for JSON serialization/deserialization
+     * @return A new instance of {@link McpSyncClient}
+     */
     public static McpSyncClient sync(McpTransport transport, Duration requestTimeout,
             ObjectMapper objectMapper) {
         return new McpSyncClient(async(transport, requestTimeout, objectMapper));
     }
-
-
 }

mcp/src/main/java/spring/ai/experimental/mcp/client/stdio/StdioServerTransport.java

@@ -10,10 +10,11 @@ import java.util.concurrent.Executors;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 import spring.ai.experimental.mcp.client.util.Assert;
-import spring.ai.experimental.mcp.spec.DefaultMcpTransport;
+import spring.ai.experimental.mcp.spec.AbstractMcpTransport;
 import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCMessage;
 import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCNotification;
 import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCRequest;
@@ -25,7 +26,7 @@ import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCResponse;
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
-public class StdioServerTransport extends DefaultMcpTransport {
+public class StdioServerTransport extends AbstractMcpTransport {
 
 	private final static TypeReference<HashMap<String, Object>> MAP_TYPE_REF = new TypeReference<HashMap<String, Object>>() {
 	};
@@ -194,47 +195,46 @@ public class StdioServerTransport extends DefaultMcpTransport {
 				.subscribe();
 	}
 
-	// TODO: provide a non-blocking variant with graceful option
-	public void stop() {
-		this.outboundScheduler.dispose();
-		this.inboundScheduler.dispose();
-		this.errorScheduler.dispose();
-
-		// Destroy process
-		if (this.process != null) {
-			this.process.destroyForcibly();
-		}
-	}
-
 	@Override
-	public void close() {
+	public Mono<Void> closeGracefully() {
 
-		this.stop();
+		this.isRunning = false;
 
-		super.close(); // Do we need this?
+		return Mono.whenDelayError(
-
+				Mono.fromRunnable(() -> {
-		// Close resources
-		if (this.processReader != null) {
-			try {
-				this.processReader.close();
-			} catch (IOException e) {
-				e.printStackTrace();
-			}
-		}
-		if (this.processWriter != null) {
-			try {
-				this.processWriter.close();
-			} catch (IOException e) {
-				e.printStackTrace();
-			}
-		}
-		if (this.processErrorReader != null) {
 					try {
 						this.processErrorReader.close();
 					} catch (IOException e) {
-				e.printStackTrace();
+						throw new RuntimeException(e);
+					}
+				}).subscribeOn(errorScheduler),
+				Mono.fromRunnable(() -> {
+					try {
+						this.processReader.close();
+					} catch (IOException e) {
+						throw new RuntimeException(e);
+					}
+				}).subscribeOn(inboundScheduler),
+				Mono.fromRunnable(() -> {
+					try {
+						this.processWriter.close();
+					} catch (IOException e) {
+						throw new RuntimeException(e);
+					}
+				}).subscribeOn(outboundScheduler))
+				.then(Mono.whenDelayError(inboundScheduler.disposeGracefully(),
+						outboundScheduler.disposeGracefully(),
+						errorScheduler.disposeGracefully()))
+				.then(Mono.fromRunnable(() -> {
+					if (this.process != null) {
+						var process = this.process.destroyForcibly();
+						if (process.exitValue() != 0) {
+							throw new RuntimeException("Failed to destroy process");
 						}
 					}
+				}))
+				.then()
+				.subscribeOn(Schedulers.boundedElastic());
 	}
 
 }

mcp/src/main/java/spring/ai/experimental/mcp/spec/AbstractMcpTransport.java

@@ -0,0 +1,119 @@
+/*
+* 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.experimental.mcp.spec;
+
+import java.util.function.Consumer;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
+import spring.ai.experimental.mcp.client.util.Assert;
+import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCMessage;
+
+/**
+ * @author Christian Tzolov
+ * @author Dariusz Jędrzejczyk
+ */
+public abstract class AbstractMcpTransport implements McpTransport {
+
+	protected final ObjectMapper objectMapper;
+
+	private final Sinks.Many<String> errorSink;
+	private final Sinks.Many<JSONRPCMessage> inboundSink;
+	private final Sinks.Many<JSONRPCMessage> outboundSink;
+
+	private Consumer<JSONRPCMessage> inboundMessageHandler = message -> System.out
+			.println("Received message: " + message);
+	private Consumer<String> errorHandler = error -> System.err.println("Received error: " + error);
+
+	public AbstractMcpTransport() {		
+		this(new ObjectMapper());
+	}
+
+	public AbstractMcpTransport(ObjectMapper objectMapper) {
+
+		Assert.notNull(objectMapper, "ObjectMapper must not be null");
+		this.objectMapper = objectMapper;
+
+		// 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.handleIncomingMessages();
+	}
+
+	public ObjectMapper getObjectMapper() {
+		return this.objectMapper;
+	}
+
+	private void handleIncomingMessages() {
+		this.inboundSink.asFlux()
+				.subscribe(message -> this.inboundMessageHandler.accept(message));
+	}
+
+	private void handleIncomingErrors() {
+		this.errorSink.asFlux().subscribe(e -> {
+			this.errorHandler.accept(e);
+			// System.err.println(e); TODO: log the error
+		});
+	}
+
+	protected Sinks.Many<JSONRPCMessage> getInboundSink() {
+		return inboundSink;
+	}
+
+	protected Sinks.Many<JSONRPCMessage> getOutboundSink() {
+		return outboundSink;
+	}
+
+	protected Sinks.Many<String> getErrorSink() {
+		return errorSink;
+	}
+
+	public void setInboudMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler) {
+		this.inboundMessageHandler = inboundMessageHandler;
+	}
+
+	public void setInboundErrorHandler(Consumer<String> errorHandler) {
+		this.errorHandler = errorHandler;
+	}
+
+	// Start processing incoming messages
+	public void start() {
+		this.handleIncomingErrors();
+		this.handleIncomingMessages();
+	}
+
+	@Override
+	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
+			// 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"));
+		}
+	}
+}

mcp/src/main/java/spring/ai/experimental/mcp/spec/McpSession.java

@@ -111,7 +111,6 @@ public class McpSession {
 		McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
 				method, params);
 		try {
-			// TODO: make it non-blocking
 			this.transport.sendMessage(jsonrpcNotification);
 		} catch (Exception e) {
 			return Mono.error(new McpError(e));

mcp/src/main/java/spring/ai/experimental/mcp/spec/McpTransport.java

@@ -15,64 +15,101 @@
  */
 package spring.ai.experimental.mcp.spec;
 
+import java.time.Duration;
 import java.util.function.Consumer;
 
 import reactor.core.publisher.Mono;
 import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCMessage;
 
 /**
- * Defines the asynchronous transport layer for the Model Context Protocol (MCP).
+ * Defines the asynchronous transport layer for the Model Context Protocol
+ * (MCP).
  * 
- * <p>The McpTransport interface provides the foundation for implementing custom transport
+ * <p>
- * mechanisms in the Model Context Protocol. It handles the bidirectional communication between
+ * The McpTransport interface provides the foundation for implementing custom
- * the client and server components, supporting asynchronous message exchange using JSON-RPC format.</p>
+ * transport
+ * mechanisms in the Model Context Protocol. It handles the bidirectional
+ * communication between
+ * the client and server components, supporting asynchronous message exchange
+ * using JSON-RPC format.
+ * </p>
  * 
- * <p>Implementations of this interface are responsible for:</p>
+ * <p>
+ * Implementations of this interface are responsible for:
+ * </p>
  * <ul>
  * <li>Managing the lifecycle of the transport connection</li>
  * <li>Handling incoming messages and errors from the server</li>
  * <li>Sending outbound messages to the server</li>
  * </ul>
  * 
- * <p>The transport layer is designed to be protocol-agnostic, allowing for various implementations
+ * <p>
- * such as WebSocket, HTTP, or custom protocols.</p>
+ * The transport layer is designed to be protocol-agnostic, allowing for various
+ * implementations
+ * such as WebSocket, HTTP, or custom protocols.
+ * </p>
  *
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
-public interface McpTransport extends AutoCloseable {
+public interface McpTransport {
 
     /**
      * Initializes and starts the transport connection.
      * 
-     * <p>This method should be called before any message exchange can occur. It sets up
+     * <p>
-     * the necessary resources and establishes the connection to the server.</p>
+     * This method should be called before any message exchange can occur. It sets
+     * up
+     * the necessary resources and establishes the connection to the server.
+     * </p>
      */
     void start();
 
     /**
      * Closes the transport connection and releases any associated resources.
      * 
-     * <p>This method ensures proper cleanup of resources when the transport is no longer needed.
+     * <p>
-     * It should handle the graceful shutdown of any active connections.</p>
+     * This method ensures proper cleanup of resources when the transport is no
+     * longer needed.
+     * It should handle the graceful shutdown of any active connections.
+     * </p>
      */
-    void close();
+    default void close() {
+        this.closeGracefully().subscribe();
+    }
+
+    /**
+     * Closes the transport connection and releases any associated resources
+     * asynchronously.
+     * 
+     * @return a {@link Mono<Void>} that completes when the connection has been
+     *         closed.
+     */
+    Mono<Void> closeGracefully();
 
     /**
      * Sets the handler for processing incoming messages from the server.
      * 
-     * <p>The provided handler will be called whenever a new message is received from the server.
+     * <p>
-     * Messages are expected to follow the JSON-RPC format as defined in the MCP specification.</p>
+     * The provided handler will be called whenever a new message is received from
+     * the server.
+     * Messages are expected to follow the JSON-RPC format as defined in the MCP
+     * specification.
+     * </p>
      *
-     * @param inboundMessageHandler a consumer that processes incoming {@link JSONRPCMessage}s
+     * @param inboundMessageHandler a consumer that processes incoming
+     *                              {@link JSONRPCMessage}s
      */
     void setInboudMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler);
 
     /**
      * Sets the handler for processing transport-level errors.
      * 
-     * <p>The provided handler will be called when errors occur during transport operations,
+     * <p>
-     * such as connection failures or protocol violations.</p>
+     * The provided handler will be called when errors occur during transport
+     * operations,
+     * such as connection failures or protocol violations.
+     * </p>
      *
      * @param inboundErrorHandler a consumer that processes error messages
      */
@@ -81,8 +118,11 @@ public interface McpTransport extends AutoCloseable {
     /**
      * Sends a message to the server asynchronously.
      * 
-     * <p>This method handles the transmission of messages to the server in an asynchronous manner.
+     * <p>
-     * Messages are sent in JSON-RPC format as specified by the MCP protocol.</p>
+     * This method handles the transmission of messages to the server in an
+     * asynchronous manner.
+     * Messages are sent in JSON-RPC format as specified by the MCP protocol.
+     * </p>
      *
      * @param message the {@link JSONRPCMessage} to be sent to the server
      * @return a {@link Mono<Void>} that completes when the message has been sent

mcp/src/test/java/spring/ai/experimental/mcp/spec/DefaultMcpTransportTests.java

@@ -15,41 +15,55 @@
  */
 package spring.ai.experimental.mcp.spec;
 
+import java.util.concurrent.atomic.AtomicReference;
+
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
-import spring.ai.experimental.mcp.spec.DefaultMcpTransport;
 import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCMessage;
 import spring.ai.experimental.mcp.spec.McpSchema.JSONRPCRequest;
 
-import java.util.concurrent.atomic.AtomicReference;
-
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 class DefaultMcpTransportTests {
 
-    private DefaultMcpTransport transport;
+	private AbstractMcpTransport transport;
 	private ObjectMapper objectMapper;
 
 	@BeforeEach
 	void setUp() {
 		objectMapper = new ObjectMapper();
-        transport = new DefaultMcpTransport(objectMapper);
+		transport = new AbstractMcpTransport(objectMapper) {
+			@Override
+			public Mono<Void> closeGracefully() {
+				return null;
+			}
+		};
 	}
 
 	@Test
 	void constructorShouldValidateObjectMapper() {
-        assertThatThrownBy(() -> new DefaultMcpTransport(null))
+		assertThatThrownBy(() -> new AbstractMcpTransport(null) {
+			@Override
+			public Mono<Void> closeGracefully() {
+				return null;
+			}
+		})
 				.isInstanceOf(IllegalArgumentException.class)
 				.hasMessage("ObjectMapper must not be null");
 	}
 
 	@Test
 	void defaultConstructorShouldCreateValidInstance() {
-        DefaultMcpTransport defaultTransport = new DefaultMcpTransport();
+		AbstractMcpTransport defaultTransport = new AbstractMcpTransport() {
+			@Override
+			public Mono<Void> closeGracefully() {
+				return null;
+			}
+		};
 		assertThat(defaultTransport.getObjectMapper()).isNotNull();
 	}