From 68c21d22647221c87396eff9409367f5ef40136b Mon Sep 17 00:00:00 2001
From: Christian Tzolov <ctzolov@vmware.com>
Date: Mon, 9 Dec 2024 23:05:23 +0100
Subject: [PATCH] Clean the code and add more tests

---
 mcp/pom.xml                                   |   7 +-
 .../spring/ai/mcp/client/McpAsyncClient.java  |   1 +
 .../java/spring/ai/mcp/client/McpClient.java  |  21 ---
 ...rParameters.java => ServerParameters.java} |  10 +-
 .../ServerParametersParser.java}              |  19 ++-
 .../client/stdio/StdioServerTransport.java    | 160 ++++++++---------
 .../ai/mcp/spec/DefaultMcpTransport.java      |  36 ++--
 .../mcp/{client => spec}/McpAsyncSession.java |   6 +-
 .../spring/ai/mcp/spec/McpAsyncTransport.java |   2 +-
 .../java/spring/ai/mcp/spec/McpSession.java   | 155 -----------------
 .../java/spring/ai/mcp/spec/McpTransport.java |  38 -----
 .../spring/ai/mcp/McpAsyncClientTests.java    |   6 +-
 .../spring/ai/mcp/McpSyncClientTests.java     |  28 ++-
 .../ai/mcp/attic/ClientSessionTests2.java     |   4 +-
 .../ai/mcp/attic/ClientSessionTests3.java     |   4 +-
 .../ai/mcp/spec/DefaultMcpTransportTests.java | 161 ++++++++++++++++++
 .../ai/mcp/spec/McpSessionJSONRPCTest.java    | 158 -----------------
 .../spring/ai/mcp/spec/McpSessionTest.java    |  93 ----------
 .../spring/ai/mcp/spec/McpTransportTest.java  | 128 --------------
 .../java/spring/ai/mcp/McpApplication.java    |   7 +-
 20 files changed, 312 insertions(+), 732 deletions(-)
 rename mcp/src/main/java/spring/ai/mcp/client/stdio/{StdioServerParameters.java => ServerParameters.java} (90%)
 rename mcp/src/main/java/spring/ai/mcp/client/{util/McpServerParameterParser.java => stdio/ServerParametersParser.java} (79%)
 rename mcp/src/main/java/spring/ai/mcp/{client => spec}/McpAsyncSession.java (95%)
 delete mode 100644 mcp/src/main/java/spring/ai/mcp/spec/McpSession.java
 delete mode 100644 mcp/src/main/java/spring/ai/mcp/spec/McpTransport.java
 create mode 100644 mcp/src/test/java/spring/ai/mcp/spec/DefaultMcpTransportTests.java
 delete mode 100644 mcp/src/test/java/spring/ai/mcp/spec/McpSessionJSONRPCTest.java
 delete mode 100644 mcp/src/test/java/spring/ai/mcp/spec/McpSessionTest.java
 delete mode 100644 mcp/src/test/java/spring/ai/mcp/spec/McpTransportTest.java

diff --git a/mcp/pom.xml b/mcp/pom.xml
index 536e103..9cb32e7 100644
--- a/mcp/pom.xml
+++ b/mcp/pom.xml
@@ -55,6 +55,11 @@
 			<version>5.11.0</version>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>io.projectreactor</groupId>
+			<artifactId>reactor-test</artifactId>
+			<scope>test</scope>
+		</dependency>
 
 		<!-- <dependency>
 			<groupId>org.springframework.ai</groupId>
@@ -114,4 +119,4 @@
 		</repository>
 	</repositories>
 
-</project>
\ No newline at end of file
+</project>
diff --git a/mcp/src/main/java/spring/ai/mcp/client/McpAsyncClient.java b/mcp/src/main/java/spring/ai/mcp/client/McpAsyncClient.java
index bdd961a..90d0225 100644
--- a/mcp/src/main/java/spring/ai/mcp/client/McpAsyncClient.java
+++ b/mcp/src/main/java/spring/ai/mcp/client/McpAsyncClient.java
@@ -5,6 +5,7 @@ import java.time.Duration;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import reactor.core.publisher.Mono;
+import spring.ai.mcp.spec.McpAsyncSession;
 import spring.ai.mcp.spec.McpAsyncTransport;
 import spring.ai.mcp.spec.McpSchema;
 
diff --git a/mcp/src/main/java/spring/ai/mcp/client/McpClient.java b/mcp/src/main/java/spring/ai/mcp/client/McpClient.java
index d85125a..06d2fb7 100644
--- a/mcp/src/main/java/spring/ai/mcp/client/McpClient.java
+++ b/mcp/src/main/java/spring/ai/mcp/client/McpClient.java
@@ -16,30 +16,9 @@
 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 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;
-import spring.ai.mcp.spec.McpSchema.CallToolRequest;
-import spring.ai.mcp.spec.McpSchema.CallToolResult;
-import spring.ai.mcp.spec.McpSchema.ClientCapabilities;
-import spring.ai.mcp.spec.McpSchema.Implementation;
-import spring.ai.mcp.spec.McpSchema.InitializeRequest;
-import spring.ai.mcp.spec.McpSchema.InitializeResult;
-import spring.ai.mcp.spec.McpSchema.ListResourceTemplatesResult;
-import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
-import spring.ai.mcp.spec.McpSchema.ListToolsResult;
-import spring.ai.mcp.spec.McpSchema.PaginatedRequest;
-import spring.ai.mcp.spec.McpSchema.ReadResourceRequest;
-import spring.ai.mcp.spec.McpSchema.Resource;
-import spring.ai.mcp.spec.McpSchema.SubscribeRequest;
-import spring.ai.mcp.spec.McpSchema.UnsubscribeRequest;
-import spring.ai.mcp.spec.McpSchema.ClientCapabilities.RootCapabilities;
 
 /**
  * The MCP client is the main entry point for interacting with the Model Context Protocol
diff --git a/mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerParameters.java b/mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParameters.java
similarity index 90%
rename from mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerParameters.java
rename to mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParameters.java
index d9c3429..a7a7851 100644
--- a/mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerParameters.java
+++ b/mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParameters.java
@@ -7,13 +7,15 @@ import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import spring.ai.mcp.client.util.Assert;
 
 /**
  * Server parameters for stdio client.
  */
-public class StdioServerParameters {
+@JsonInclude(JsonInclude.Include.NON_ABSENT)
+public class ServerParameters {
 
 	// Environment variables to inherit by default
 	private static final List<String> DEFAULT_INHERITED_ENV_VARS = System.getProperty("os.name")
@@ -32,7 +34,7 @@ public class StdioServerParameters {
 	@JsonProperty("env")
 	private Map<String, String> env;
 
-	private StdioServerParameters(String command, List<String> args, Map<String, String> env) {
+	private ServerParameters(String command, List<String> args, Map<String, String> env) {
 		Assert.notNull(command, "The command can not be null");
 		Assert.notNull(args, "The args can not be null");
 
@@ -105,8 +107,8 @@ public class StdioServerParameters {
 			return this;
 		}
 
-		public StdioServerParameters build() {
-			return new StdioServerParameters(command, args, env);
+		public ServerParameters build() {
+			return new ServerParameters(command, args, env);
 		}
 
 	}
diff --git a/mcp/src/main/java/spring/ai/mcp/client/util/McpServerParameterParser.java b/mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParametersParser.java
similarity index 79%
rename from mcp/src/main/java/spring/ai/mcp/client/util/McpServerParameterParser.java
rename to mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParametersParser.java
index 0dc61b5..a5c407b 100644
--- a/mcp/src/main/java/spring/ai/mcp/client/util/McpServerParameterParser.java
+++ b/mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParametersParser.java
@@ -13,30 +13,31 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-package spring.ai.mcp.client.util;
+package spring.ai.mcp.client.stdio;
 
 import java.util.Map;
+import java.util.List;
 
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import spring.ai.mcp.client.util.McpServerParameterParser.McpServers.McpServer;
+import spring.ai.mcp.client.stdio.ServerParametersParser.McpServerConfigurations.McpServerConfiguration;
 
 /**
  * @author Christian Tzolov
  * @since 1.0.0
  */
 
-public class McpServerParameterParser {
+public class ServerParametersParser {
 
 	@JsonInclude(JsonInclude.Include.NON_ABSENT)
-	public record McpServers(// @formatter:off
-		@JsonProperty("mcpServers") Map<String, McpServer> mcpServers) {
+	public record McpServerConfigurations(// @formatter:off
+		@JsonProperty("mcpServers") Map<String, McpServerConfiguration> mcpServers) {
 
 		@JsonInclude(JsonInclude.Include.NON_ABSENT)
-		public record McpServer(
+		public record McpServerConfiguration(
 			@JsonProperty("command") String command,
-			@JsonProperty("args") String[] args,
+			@JsonProperty("args") List<String> args,
 			@JsonProperty("env") Map<String, String> env) {
 		}
 	} // @formatter:on
@@ -72,10 +73,10 @@ public class McpServerParameterParser {
 								""";
 
 		// Deserialize JSON into McpServers
-		McpServers servers = objectMapper.readValue(jsonInput, McpServers.class);
+		McpServerConfigurations servers = objectMapper.readValue(jsonInput, McpServerConfigurations.class);
 
 		// Access individual servers
-		for (Map.Entry<String, McpServer> entry : servers.mcpServers().entrySet()) {
+		for (Map.Entry<String, McpServerConfiguration> entry : servers.mcpServers().entrySet()) {
 			System.out.println("Key: " + entry.getKey());
 			System.out.println("Command: " + entry.getValue().command());
 			System.out.println("Arguments: " + String.join(", ", entry.getValue().args()));
diff --git a/mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java b/mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java
index 3ff5e20..35f80c2 100644
--- a/mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java
+++ b/mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java
@@ -3,12 +3,9 @@ package spring.ai.mcp.client.stdio;
 import java.io.BufferedReader;
 import java.io.BufferedWriter;
 import java.io.IOException;
-import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.Objects;
 import java.util.concurrent.Executors;
 
 import com.fasterxml.jackson.core.type.TypeReference;
@@ -27,12 +24,13 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
  */
 public class StdioServerTransport extends DefaultMcpTransport {
 
+	private final static TypeReference<HashMap<String, Object>> MAP_TYPE_REF = new TypeReference<HashMap<String, Object>>() {
+	};
+
 	private Process process;
 
 	private BufferedReader processErrorReader;
-
 	private BufferedReader processReader;
-
 	private BufferedWriter processWriter;
 
 	private Scheduler inboundScheduler;
@@ -41,29 +39,28 @@ public class StdioServerTransport extends DefaultMcpTransport {
 
 	private volatile boolean isRunning;
 
-	private final StdioServerParameters params;
+	private final ServerParameters params;
 
-	public StdioServerTransport(StdioServerParameters params) {
-		this(params, Duration.ofMillis(100));
+	public StdioServerTransport(ServerParameters params) {
+		this(params, new ObjectMapper());
 	}
 
-	public StdioServerTransport(StdioServerParameters params, Duration writeTimeout) {
-		super(writeTimeout);
+	public StdioServerTransport(ServerParameters params, ObjectMapper objectMapper) {
+
+		super(objectMapper);
 
 		Assert.notNull(params, "The params can not be null");
+		Assert.notNull(objectMapper, "The ObjectMapper can not be null");
 
 		this.params = params;
 
 		// Start threads
-		this.inboundScheduler =
-				Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
-						"inbound");
-		this.outboundScheduler =
-				Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
-						"outbound");
-		this.errorScheduler =
-				Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
-						"error");
+		this.inboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
+				"inbound");
+		this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
+				"outbound");
+		this.errorScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
+				"error");
 	}
 
 	@Override
@@ -79,8 +76,7 @@ public class StdioServerTransport extends DefaultMcpTransport {
 		// Start the process
 		try {
 			this.process = processBuilder.start();
-		}
-		catch (IOException e) {
+		} catch (IOException e) {
 			throw new RuntimeException("Failed to start process with command: " + fullCommand, e);
 		}
 
@@ -100,14 +96,12 @@ public class StdioServerTransport extends DefaultMcpTransport {
 		startInboundProcessing();
 		startOutboundProcessing();
 		startErrorProcessing();
-
 	}
 
 	public void awaitForExit() {
 		try {
 			this.process.waitFor();
-		}
-		catch (InterruptedException e) {
+		} catch (InterruptedException e) {
 			throw new RuntimeException("Process interrupted", e);
 		}
 	}
@@ -121,56 +115,20 @@ public class StdioServerTransport extends DefaultMcpTransport {
 						System.out.println("Received error line: " + line);
 						// TODO: handle errors, etc.
 						this.getErrorSink().tryEmitNext(line);
-					}
-					catch (Exception e) {
+					} catch (Exception e) {
 						throw new RuntimeException(e);
 					}
 				}
-			}
-			catch (IOException e) {
+			} catch (IOException e) {
 				if (this.isRunning) {
 					throw new RuntimeException(e);
 				}
-			}
-			finally {
+			} finally {
 				this.isRunning = false;
 			}
 		});
 	}
 
-	private static ObjectMapper OBJECT_MAPPER = new ObjectMapper();
-
-	private static TypeReference<HashMap<String, Object>> MAP_TYPE_REF = new TypeReference<HashMap<String, Object>>() {
-
-	};
-
-	public static Map<String, Object> jsonToMap(String json) {
-		try {
-			return OBJECT_MAPPER.readValue(json, MAP_TYPE_REF);
-		}
-		catch (Exception e) {
-			throw new RuntimeException(e);
-		}
-	}
-
-	private JSONRPCMessage deserializeJsonRpcMessage(String jsonText) throws IOException {
-
-		var map = jsonToMap(jsonText);
-
-		// Determine message type based on specific JSON structure
-		if (map.containsKey("method") && map.containsKey("id")) {
-			return OBJECT_MAPPER.convertValue(map, JSONRPCRequest.class);
-		}
-		else if (map.containsKey("method") && !map.containsKey("id")) {
-			return OBJECT_MAPPER.convertValue(map, JSONRPCNotification.class);
-		}
-		else if (map.containsKey("result") || map.containsKey("error")) {
-			return OBJECT_MAPPER.convertValue(map, JSONRPCResponse.class);
-		}
-
-		throw new IllegalArgumentException("Cannot deserialize JSONRPCMessage: " + jsonText);
-	}
-
 	private void startInboundProcessing() {
 		this.inboundScheduler.schedule(() -> {
 			try {
@@ -182,43 +140,55 @@ public class StdioServerTransport extends DefaultMcpTransport {
 							// TODO: Back off, reschedule, give up?
 							throw new RuntimeException("Failed to enqueue message");
 						}
-					}
-					catch (Exception e) {
+					} catch (Exception e) {
 						throw new RuntimeException(e);
 					}
 				}
-			}
-			catch (IOException e) {
+			} catch (IOException e) {
 				if (isRunning) {
 					throw new RuntimeException(e);
 				}
-			}
-			finally {
+			} finally {
 				isRunning = false;
 			}
 		});
 	}
 
+	private JSONRPCMessage deserializeJsonRpcMessage(String jsonText) throws IOException {
+
+		var map = this.objectMapper.readValue(jsonText, MAP_TYPE_REF);
+
+		// Determine message type based on specific JSON structure
+		if (map.containsKey("method") && map.containsKey("id")) {
+			return this.objectMapper.convertValue(map, JSONRPCRequest.class);
+		} else if (map.containsKey("method") && !map.containsKey("id")) {
+			return this.objectMapper.convertValue(map, JSONRPCNotification.class);
+		} else if (map.containsKey("result") || map.containsKey("error")) {
+			return this.objectMapper.convertValue(map, JSONRPCResponse.class);
+		}
+
+		throw new IllegalArgumentException("Cannot deserialize JSONRPCMessage: " + jsonText);
+	}
+
 	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(OBJECT_MAPPER.writeValueAsString(message));
-					    this.processWriter.newLine();
-					    this.processWriter.flush();
-					    s.next(message);
-				    }
-				    catch (IOException e) {
-					    s.error(new RuntimeException(e));
-				    }
-			    }
-		    })
-		    .subscribe();
+				.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(objectMapper.writeValueAsString(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
@@ -235,24 +205,30 @@ public class StdioServerTransport extends DefaultMcpTransport {
 
 	@Override
 	public void close() {
-		stop();
-		
+
+		this.stop();
+
 		super.close(); // Do we need this?
 
 		// Close resources
 		if (this.processReader != null) {
 			try {
 				this.processReader.close();
-			}
-			catch (IOException e) {
+			} catch (IOException e) {
 				e.printStackTrace();
 			}
 		}
 		if (this.processWriter != null) {
 			try {
 				this.processWriter.close();
+			} catch (IOException e) {
+				e.printStackTrace();
 			}
-			catch (IOException e) {
+		}
+		if (this.processErrorReader != null) {
+			try {
+				this.processErrorReader.close();
+			} catch (IOException e) {
 				e.printStackTrace();
 			}
 		}
diff --git a/mcp/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java b/mcp/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java
index 9b3a2d9..5631bd8 100644
--- a/mcp/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java
+++ b/mcp/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java
@@ -15,9 +15,9 @@
 */
 package spring.ai.mcp.spec;
 
-import java.time.Duration;
 import java.util.function.Consumer;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
 import spring.ai.mcp.client.util.Assert;
@@ -29,19 +29,24 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
  */
 public class DefaultMcpTransport implements McpAsyncTransport {
 
+	protected final ObjectMapper objectMapper;
+
 	private final Sinks.Many<String> errorSink;
-
 	private final Sinks.Many<JSONRPCMessage> inboundSink;
-
 	private final Sinks.Many<JSONRPCMessage> outboundSink;
 
-	private final Duration writeTimeout;
+	private Consumer<JSONRPCMessage> inboundMessageHandler = message -> System.out
+			.println("Received message: " + message);
+	private Consumer<String> errorHandler = error -> System.err.println("Received error: " + error);
 
-	private Consumer<JSONRPCMessage> messageReader = message -> System.out.println("Received message: " + message);
+	public DefaultMcpTransport() {		
+		this(new ObjectMapper());
+	}
 
-	private Consumer<String> errorReader = error -> System.err.println("Received error: " + error);
+	public DefaultMcpTransport(ObjectMapper objectMapper) {
 
-	public DefaultMcpTransport(Duration readTimeout) {
+		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.
@@ -52,19 +57,22 @@ public class DefaultMcpTransport implements McpAsyncTransport {
 		this.errorSink = Sinks.many().unicast().onBackpressureBuffer();
 		this.inboundSink = Sinks.many().unicast().onBackpressureBuffer();
 		this.outboundSink = Sinks.many().unicast().onBackpressureBuffer();
-		this.writeTimeout = readTimeout;
 
 		this.handleIncomingMessages();
 	}
 
+	public ObjectMapper getObjectMapper() {
+		return this.objectMapper;
+	}
+
 	private void handleIncomingMessages() {
 		this.inboundSink.asFlux()
-				.subscribe(message -> this.messageReader.accept(message));
+				.subscribe(message -> this.inboundMessageHandler.accept(message));
 	}
 
 	private void handleIncomingErrors() {
 		this.errorSink.asFlux().subscribe(e -> {
-			this.errorReader.accept(e);
+			this.errorHandler.accept(e);
 			System.err.println(e);
 		});
 	}
@@ -81,12 +89,12 @@ public class DefaultMcpTransport implements McpAsyncTransport {
 		return errorSink;
 	}
 
-	public void setMessageHandler(Consumer<JSONRPCMessage> messageReader) {
-		this.messageReader = messageReader;
+	public void setInboudMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler) {
+		this.inboundMessageHandler = inboundMessageHandler;
 	}
 
-	public void setErrorHandler(Consumer<String> errorReader) {
-		this.errorReader = errorReader;
+	public void setErrorHandler(Consumer<String> errorHandler) {
+		this.errorHandler = errorHandler;
 	}
 
 	// Start processing incoming messages
diff --git a/mcp/src/main/java/spring/ai/mcp/client/McpAsyncSession.java b/mcp/src/main/java/spring/ai/mcp/spec/McpAsyncSession.java
similarity index 95%
rename from mcp/src/main/java/spring/ai/mcp/client/McpAsyncSession.java
rename to mcp/src/main/java/spring/ai/mcp/spec/McpAsyncSession.java
index 11b775f..42560bc 100644
--- a/mcp/src/main/java/spring/ai/mcp/client/McpAsyncSession.java
+++ b/mcp/src/main/java/spring/ai/mcp/spec/McpAsyncSession.java
@@ -1,4 +1,4 @@
-package spring.ai.mcp.client;
+package spring.ai.mcp.spec;
 
 import java.time.Duration;
 import java.util.Map;
@@ -10,8 +10,6 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoSink;
 import spring.ai.mcp.client.util.Assert;
-import spring.ai.mcp.spec.McpAsyncTransport;
-import spring.ai.mcp.spec.McpSchema;
 
 public class McpAsyncSession {
 
@@ -33,7 +31,7 @@ public class McpAsyncSession {
 		this.objectMapper = objectMapper;
 		this.transport = transport;
 
-		this.transport.setMessageHandler(message -> {
+		this.transport.setInboudMessageHandler(message -> {
 			switch (message) {
 				case McpSchema.JSONRPCResponse response -> {
 					var sink = pendingResponses.remove(response.id());
diff --git a/mcp/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java b/mcp/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java
index b8109c7..264d4a6 100644
--- a/mcp/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java
+++ b/mcp/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java
@@ -31,7 +31,7 @@ public interface McpAsyncTransport extends AutoCloseable{
 
 	void close();
 
-    void setMessageHandler(Consumer<JSONRPCMessage> messageReader);
+    void setInboudMessageHandler(Consumer<JSONRPCMessage> messageReader);
 
 	void setErrorHandler(Consumer<String> errorReader);
 
diff --git a/mcp/src/main/java/spring/ai/mcp/spec/McpSession.java b/mcp/src/main/java/spring/ai/mcp/spec/McpSession.java
deleted file mode 100644
index 7fa2bc4..0000000
--- a/mcp/src/main/java/spring/ai/mcp/spec/McpSession.java
+++ /dev/null
@@ -1,155 +0,0 @@
-/*
-* 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.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 spring.ai.mcp.client.util.Assert;
-import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
-import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
-import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
-
-public class McpSession implements AutoCloseable {
-
-	final ConcurrentHashMap<Object, CompletableFuture<JSONRPCResponse>> waitingForResponseStreams;
-
-	private final Duration readTimeout;
-
-	private final ObjectMapper objectMapper;
-
-	private final McpTransport transport;
-
-	public McpSession(McpTransport transport, Duration readTimeout, ObjectMapper objectMapper) {
-
-		Assert.notNull(transport, "The transport can not be null");
-		Assert.notNull(readTimeout, "The readTimeout can not be null");
-		Assert.notNull(objectMapper, "The objectMapper can not be null");
-
-		this.waitingForResponseStreams = new ConcurrentHashMap<>();
-
-		this.readTimeout = readTimeout;
-		this.objectMapper = objectMapper;
-
-		this.transport = transport;
-
-		this.transport.setMessageHandler(message -> {
-			if (message instanceof JSONRPCResponse jsonRpcResponse) {
-				CompletableFuture<JSONRPCResponse> responseFuture = this.waitingForResponseStreams
-						.remove(jsonRpcResponse.id());
-				if (responseFuture != null) {
-					responseFuture.complete(jsonRpcResponse);
-				}
-			} else if (message instanceof JSONRPCRequest jsonRpcRequest) {
-				handleRequest(new RequestResponder<>(jsonRpcRequest.id(), jsonRpcRequest));
-			} else if (message instanceof JSONRPCNotification jsonRpcNotification) {
-				handleNotification(jsonRpcNotification);
-			}
-		});
-
-		this.transport.setErrorHandler(error -> {
-			System.out.println("Received error: " + error);
-		});
-
-		this.transport.start();
-	}
-
-	public <T> CompletableFuture<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
-
-		String requestId = UUID.randomUUID().toString();
-		CompletableFuture<JSONRPCResponse> responseFuture = new CompletableFuture<>();
-		this.waitingForResponseStreams.put(requestId, responseFuture);
-
-		JSONRPCRequest jsonrpcRequest = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, requestId, requestParams);
-
-		try {
-			this.transport.sendMessage(jsonrpcRequest);
-		} catch (Exception e) {
-			return CompletableFuture.failedFuture(new McpError(e));
-		}
-
-		return responseFuture.orTimeout(this.readTimeout.toMillis(), TimeUnit.MILLISECONDS)
-				.thenApply(jsonRpcResponse -> {
-					if (jsonRpcResponse.error() != null) {
-						throw new McpError(jsonRpcResponse.error());
-					}
-
-					return this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef);
-				});
-	}
-
-	public CompletableFuture<Void> sendNotification(String method) {
-		return sendNotification(method, null);
-	}
-
-	public CompletableFuture<Void> sendNotification(String method, Map<String, Object> params) {
-		JSONRPCNotification jsonrpcNotification = new JSONRPCNotification(McpSchema.JSONRPC_VERSION, method, params);
-		try {
-			this.transport.sendMessage(jsonrpcNotification);
-		} catch (Exception e) {
-			return CompletableFuture.failedFuture(new McpError(e));
-		}
-		return CompletableFuture.completedFuture(null);
-	}
-
-	protected void handleRequest(RequestResponder<JSONRPCRequest> responder) {
-		// Override to handle requests
-		System.out.println("Handling request: " + responder);
-	}
-
-	protected void handleNotification(JSONRPCNotification notification) {
-		// Override to handle notifications
-		System.out.println("Handling notification: " + notification);
-	}
-
-	@Override
-	public void close() {
-		this.transport.close();
-	}
-
-	public static class RequestResponder<R> {
-
-		private final Object requestId;
-
-		private final R request;
-
-		public RequestResponder(Object requestId, R request) {
-			this.requestId = requestId;
-			this.request = request;
-		}
-
-		public CompletableFuture<Void> respond(Object response) {
-			// Send response
-			return CompletableFuture.completedFuture(null);
-		}
-
-	}
-
-	public static class McpError extends RuntimeException {
-
-		public McpError(Object error) {
-			super(error.toString());
-		}
-
-	}
-
-}
diff --git a/mcp/src/main/java/spring/ai/mcp/spec/McpTransport.java b/mcp/src/main/java/spring/ai/mcp/spec/McpTransport.java
deleted file mode 100644
index 819ef5c..0000000
--- a/mcp/src/main/java/spring/ai/mcp/spec/McpTransport.java
+++ /dev/null
@@ -1,38 +0,0 @@
-/* 
-* 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 spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
-
-/**
- * @author Christian Tzolov
- * @since 1.0.0
- */
-
-public interface McpTransport extends AutoCloseable{
-    
-    void start();
-
-	void close();
-
-    void setMessageHandler(Consumer<JSONRPCMessage> messageReader);
-
-	void setErrorHandler(Consumer<String> errorReader);
-
-	void sendMessage(JSONRPCMessage message);    
-}
diff --git a/mcp/src/test/java/spring/ai/mcp/McpAsyncClientTests.java b/mcp/src/test/java/spring/ai/mcp/McpAsyncClientTests.java
index 4629477..c87ae6b 100644
--- a/mcp/src/test/java/spring/ai/mcp/McpAsyncClientTests.java
+++ b/mcp/src/test/java/spring/ai/mcp/McpAsyncClientTests.java
@@ -25,7 +25,7 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import spring.ai.mcp.client.McpAsyncClient;
 import spring.ai.mcp.client.McpClient;
-import spring.ai.mcp.client.stdio.StdioServerParameters;
+import spring.ai.mcp.client.stdio.ServerParameters;
 import spring.ai.mcp.client.stdio.StdioServerTransport;
 import spring.ai.mcp.spec.McpSchema.CallToolRequest;
 import spring.ai.mcp.spec.McpSchema.Resource;
@@ -45,7 +45,7 @@ class McpAsyncClientTests {
 
 	private McpAsyncClient mcpAsyncClient;
 
-	private StdioServerParameters stdioParams;
+	private ServerParameters stdioParams;
 
 	private static final Duration TIMEOUT = Duration.ofSeconds(10);
 
@@ -53,7 +53,7 @@ class McpAsyncClientTests {
 
 	@BeforeEach
 	void setUp() {
-		stdioParams = StdioServerParameters.builder("npx")
+		stdioParams = ServerParameters.builder("npx")
 			.args("-y", "@modelcontextprotocol/server-everything", "dir")
 			.build();
 
diff --git a/mcp/src/test/java/spring/ai/mcp/McpSyncClientTests.java b/mcp/src/test/java/spring/ai/mcp/McpSyncClientTests.java
index b4f049b..16a4f52 100644
--- a/mcp/src/test/java/spring/ai/mcp/McpSyncClientTests.java
+++ b/mcp/src/test/java/spring/ai/mcp/McpSyncClientTests.java
@@ -16,6 +16,7 @@
 package spring.ai.mcp;
 
 import java.time.Duration;
+import java.util.List;
 import java.util.Map;
 
 import org.junit.jupiter.api.AfterEach;
@@ -28,14 +29,16 @@ import static org.assertj.core.api.Assertions.*;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import spring.ai.mcp.client.McpClient;
 import spring.ai.mcp.client.McpSyncClient;
-import spring.ai.mcp.client.stdio.StdioServerParameters;
+import spring.ai.mcp.client.stdio.ServerParameters;
 import spring.ai.mcp.client.stdio.StdioServerTransport;
 import spring.ai.mcp.spec.McpSchema.CallToolRequest;
 import spring.ai.mcp.spec.McpSchema.CallToolResult;
+import spring.ai.mcp.spec.McpSchema.Content;
 import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
 import spring.ai.mcp.spec.McpSchema.ListToolsResult;
 import spring.ai.mcp.spec.McpSchema.Tool;
 import spring.ai.mcp.spec.McpSchema.Resource;
+import spring.ai.mcp.spec.McpSchema.TextContent;
 
 /**
  * Unit tests for MCP Client Session functionality.
@@ -47,7 +50,7 @@ class McpSyncClientTests {
 
 	private McpSyncClient mcpSyncClient;
 
-	private StdioServerParameters stdioParams;
+	private ServerParameters stdioParams;
 
 	private static final Duration TIMEOUT = Duration.ofSeconds(10);
 
@@ -55,7 +58,7 @@ class McpSyncClientTests {
 
 	@BeforeEach
 	void setUp() {
-		stdioParams = StdioServerParameters.builder("npx")
+		stdioParams = ServerParameters.builder("npx")
 			.args("-y", "@modelcontextprotocol/server-everything", "dir")
 			.build();
 
@@ -98,6 +101,25 @@ class McpSyncClientTests {
 		});
 	}
 
+	@Test
+	@Timeout(15) // Giving extra time beyond the client timeout
+	void testCallTools() {
+		CallToolResult toolResult = mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4)));
+
+		assertThat(toolResult).isNotNull().satisfies(result -> {
+			
+			assertThat(result.content()).hasSize(1);
+
+			TextContent content = (TextContent) result.content().get(0);
+
+			assertThat(content).isNotNull();
+			assertThat(content.text()).isNotNull();
+			assertThat(content.text()).contains("7");
+
+			System.out.println(content);
+		});
+	}
+
 	@Test
 	@Timeout(15)
 	void testPing() {
diff --git a/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests2.java b/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests2.java
index 578fa3d..68f9eb5 100644
--- a/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests2.java
+++ b/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests2.java
@@ -20,7 +20,7 @@ import java.time.Duration;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import spring.ai.mcp.client.McpClient;
 import spring.ai.mcp.client.McpSyncClient;
-import spring.ai.mcp.client.stdio.StdioServerParameters;
+import spring.ai.mcp.client.stdio.ServerParameters;
 import spring.ai.mcp.client.stdio.StdioServerTransport;
 import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
 import spring.ai.mcp.spec.McpSchema.ListToolsResult;
@@ -34,7 +34,7 @@ public class ClientSessionTests2 {
 
 	public static void main(String[] args) {
 
-		var stdioParams = StdioServerParameters.builder("uv")
+		var stdioParams = ServerParameters.builder("uv")
 			.args("--directory", "dir", "run", "mcp-server-sqlite",
 					"--db-path", "~/test.db")
 			.build();
diff --git a/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests3.java b/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests3.java
index e080dd6..9b82709 100644
--- a/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests3.java
+++ b/mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests3.java
@@ -20,7 +20,7 @@ import java.time.Duration;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import spring.ai.mcp.client.McpClient;
 import spring.ai.mcp.client.McpSyncClient;
-import spring.ai.mcp.client.stdio.StdioServerParameters;
+import spring.ai.mcp.client.stdio.ServerParameters;
 import spring.ai.mcp.client.stdio.StdioServerTransport;
 import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
 import spring.ai.mcp.spec.McpSchema.ListToolsResult;
@@ -34,7 +34,7 @@ public class ClientSessionTests3 {
 
 	public static void main(String[] args) {
 
-		var stdioParams = StdioServerParameters.builder("uv")
+		var stdioParams = ServerParameters.builder("uv")
 			.args("--directory", "dir", "run", "mcp-server-sqlite",
 					"--db-path", "~/test.db")
 			.build();
diff --git a/mcp/src/test/java/spring/ai/mcp/spec/DefaultMcpTransportTests.java b/mcp/src/test/java/spring/ai/mcp/spec/DefaultMcpTransportTests.java
new file mode 100644
index 0000000..06ef9c8
--- /dev/null
+++ b/mcp/src/test/java/spring/ai/mcp/spec/DefaultMcpTransportTests.java
@@ -0,0 +1,161 @@
+/*
+ * 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 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.mcp.spec.McpSchema.JSONRPCMessage;
+import spring.ai.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 ObjectMapper objectMapper;
+
+    @BeforeEach
+    void setUp() {
+        objectMapper = new ObjectMapper();
+        transport = new DefaultMcpTransport(objectMapper);
+    }
+
+    @Test
+    void constructorShouldValidateObjectMapper() {
+        assertThatThrownBy(() -> new DefaultMcpTransport(null))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessage("ObjectMapper must not be null");
+    }
+
+    @Test
+    void defaultConstructorShouldCreateValidInstance() {
+        DefaultMcpTransport defaultTransport = new DefaultMcpTransport();
+        assertThat(defaultTransport.getObjectMapper()).isNotNull();
+    }
+
+    @Test
+    void customObjectMapperConstructorShouldCreateValidInstance() {
+        assertThat(transport.getObjectMapper()).isSameAs(objectMapper);
+    }
+
+    @Test
+    void sendMessageShouldEmitToOutboundSink() {
+        JSONRPCRequest message = new JSONRPCRequest(
+                "2.0",
+                "test",
+                1,
+                null);
+
+        AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
+        transport.getOutboundSink().asFlux().subscribe(receivedMessage::set);
+
+        Mono<Void> result = transport.sendMessage(message);
+
+        StepVerifier.create(result)
+                .verifyComplete();
+
+        assertThat(receivedMessage.get())
+                .isNotNull()
+                .satisfies(msg -> {
+                    assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0");
+                    assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test");
+                });
+    }
+
+    @Test
+    void customInboundMessageHandlerShouldReceiveMessages() {
+        AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
+        transport.setInboudMessageHandler(receivedMessage::set);
+        transport.start();
+
+        JSONRPCRequest message = new JSONRPCRequest(
+                "2.0",
+                "test",
+                1,
+                null);
+
+        transport.getInboundSink().tryEmitNext(message);
+
+        assertThat(receivedMessage.get())
+                .isNotNull()
+                .satisfies(msg -> {
+                    assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0");
+                    assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test");
+                });
+    }
+
+    @Test
+    void customErrorHandlerShouldReceiveErrors() {
+        AtomicReference<String> receivedError = new AtomicReference<>();
+        transport.setErrorHandler(receivedError::set);
+        transport.start();
+
+        String errorMessage = "Test error";
+        transport.getErrorSink().tryEmitNext(errorMessage);
+
+        assertThat(receivedError.get())
+                .isNotNull()
+                .isEqualTo(errorMessage);
+    }
+
+    @Test
+    void startShouldInitializeMessageAndErrorHandling() {
+        AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
+        AtomicReference<String> receivedError = new AtomicReference<>();
+
+        transport.setInboudMessageHandler(receivedMessage::set);
+        transport.setErrorHandler(receivedError::set);
+        transport.start();
+
+        JSONRPCRequest message = new JSONRPCRequest(
+                "2.0",
+                "test",
+                1,
+                null);
+
+        String errorMessage = "Test error";
+
+        transport.getInboundSink().tryEmitNext(message);
+        transport.getErrorSink().tryEmitNext(errorMessage);
+
+        assertThat(receivedMessage.get()).isNotNull();
+        assertThat(receivedError.get()).isEqualTo(errorMessage);
+    }
+
+    @Test
+    void sendMessageShouldReturnErrorWhenSinkFails() {
+        JSONRPCRequest message = new JSONRPCRequest(
+                "2.0",
+                "test",
+                1,
+                null);
+        // Close the sink to simulate failure
+        transport.getOutboundSink().tryEmitComplete();
+
+        Mono<Void> result = transport.sendMessage(message);
+
+        StepVerifier.create(result)
+                .verifyErrorSatisfies(error -> assertThat(error)
+                        .isInstanceOf(RuntimeException.class)
+                        .hasMessage("Failed to enqueue message"));
+    }
+}
diff --git a/mcp/src/test/java/spring/ai/mcp/spec/McpSessionJSONRPCTest.java b/mcp/src/test/java/spring/ai/mcp/spec/McpSessionJSONRPCTest.java
deleted file mode 100644
index 3a56a89..0000000
--- a/mcp/src/test/java/spring/ai/mcp/spec/McpSessionJSONRPCTest.java
+++ /dev/null
@@ -1,158 +0,0 @@
-/* 
-* 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 com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
-import org.junit.jupiter.api.Test;
-import org.mockito.ArgumentCaptor;
-import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
-import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
-import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
-import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
-
-import java.time.Duration;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-import java.util.function.Consumer;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.Mockito.*;
-
-class McpSessionJSONRPCTest {
-
-    private McpTransport mockTransport;
-    private ObjectMapper objectMapper;
-    private McpSession mcpSession;
-    private ArgumentCaptor<Consumer<JSONRPCMessage>> messageHandlerCaptor;
-
-    @BeforeEach
-    void setUp() {
-        mockTransport = mock(McpTransport.class);
-        objectMapper = new ObjectMapper();
-        mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), objectMapper);
-
-        // Capture the message handler set during the session initialization
-        messageHandlerCaptor = ArgumentCaptor.forClass(Consumer.class);
-        verify(mockTransport).setMessageHandler(messageHandlerCaptor.capture());
-    }
-
-    @Test
-    void sendRequestShouldSendValidRequestMessage() {
-        String method = "testMethod";
-        Map<String, Object> params = Map.of("key", "value");
-
-        // Call sendRequest
-        CompletableFuture<String> future = mcpSession.sendRequest(method, params, new TypeReference<>() {});
-
-        // Capture the JSONRPCRequest sent to the transport
-        ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
-        verify(mockTransport).sendMessage(messageCaptor.capture());
-
-        // Assert the message details
-        JSONRPCRequest request = (JSONRPCRequest) messageCaptor.getValue();
-        assertEquals(method, request.method());
-        assertEquals("value", ((Map<?, ?>) request.params()).get("key"));
-        assertNotNull(request.id());
-    }
-
-    @Test
-    void sendNotificationShouldSendValidNotificationMessage() {
-        String method = "notifyMethod";
-        Map<String, Object> params = Map.of("param1", "value1");
-
-        // Call sendNotification
-        mcpSession.sendNotification(method, params);
-
-        // Capture the JSONRPCNotification sent to the transport
-        ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
-        verify(mockTransport).sendMessage(messageCaptor.capture());
-
-        // Assert the notification details
-        JSONRPCNotification notification = (JSONRPCNotification) messageCaptor.getValue();
-        assertEquals(method, notification.method());
-        assertEquals("value1", ((Map<?, ?>) notification.params()).get("param1"));
-    }
-
-    @Test
-    void messageHandlerShouldProcessResponseCorrectly() {
-        // Retrieve the captured message handler
-        Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
-
-        // Simulate a JSONRPCResponse
-        JSONRPCResponse response = new JSONRPCResponse("2.0", "1", "resultValue", null);
-        CompletableFuture<JSONRPCResponse> responseFuture = new CompletableFuture<>();
-        mcpSession.waitingForResponseStreams.put("1", responseFuture);
-
-        // Trigger the handler with the response
-        messageHandler.accept(response);
-
-        // Verify the response is completed
-        assertTrue(responseFuture.isDone());
-        assertEquals("resultValue", responseFuture.join().result());
-    }
-
-    @Test
-    @Disabled
-    void messageHandlerShouldInvokeHandleRequestForRequestMessage() {
-        // Retrieve the captured message handler
-        Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
-
-        // Mock request data
-        JSONRPCRequest request = new JSONRPCRequest("2.0", "testMethod", "1", Map.of("key", "value"));
-
-        // Override the handleRequest to check invocation
-        McpSession spySession = spy(mcpSession);
-        doNothing().when(spySession).handleRequest(any());
-
-        // Trigger the handler with a request
-        messageHandler.accept(request);
-
-        // Verify handleRequest was called
-        verify(spySession).handleRequest(any());
-    }
-
-    @Test
-    @Disabled
-    void messageHandlerShouldInvokeHandleNotificationForNotificationMessage() {
-        // Retrieve the captured message handler
-        Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
-
-        // Mock notification data
-        JSONRPCNotification notification = new JSONRPCNotification("2.0", "notifyMethod", Map.of("key", "value"));
-
-        // Override the handleNotification to check invocation
-        McpSession spySession = spy(mcpSession);
-        doNothing().when(spySession).handleNotification(any());
-
-        // Trigger the handler with a notification
-        messageHandler.accept(notification);
-
-        // Verify handleNotification was called
-        verify(spySession).handleNotification(notification);
-    }
-
-    @Test
-    void closeShouldInvokeTransportClose() {
-        // Call close on the session
-        mcpSession.close();
-
-        // Verify transport close was called
-        verify(mockTransport).close();
-    }
-}
diff --git a/mcp/src/test/java/spring/ai/mcp/spec/McpSessionTest.java b/mcp/src/test/java/spring/ai/mcp/spec/McpSessionTest.java
deleted file mode 100644
index 2826b60..0000000
--- a/mcp/src/test/java/spring/ai/mcp/spec/McpSessionTest.java
+++ /dev/null
@@ -1,93 +0,0 @@
-/*
-* 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.time.Duration;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
-import org.junit.jupiter.api.Test;
-import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-class McpSessionTest {
-
-	private McpTransport mockTransport;
-
-	private ObjectMapper mockObjectMapper;
-
-	private McpSession mcpSession;
-
-	@BeforeEach
-	void setUp() {
-		mockTransport = mock(McpTransport.class);
-		mockObjectMapper = mock(ObjectMapper.class);
-		mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), mockObjectMapper);
-	}
-
-	@Test
-	void constructorShouldThrowIfArgumentsAreNull() {
-		assertThrows(IllegalArgumentException.class,
-				() -> new McpSession(null, Duration.ofSeconds(5), mockObjectMapper));
-		assertThrows(IllegalArgumentException.class, () -> new McpSession(mockTransport, null, mockObjectMapper));
-		assertThrows(IllegalArgumentException.class, () -> new McpSession(mockTransport, Duration.ofSeconds(5), null));
-	}
-
-	@Test
-	@Disabled
-	void sendRequestShouldSendJsonRpcRequest() throws Exception {
-		String method = "testMethod";
-		Object params = Map.of("key", "value");
-		TypeReference<String> typeRef = new TypeReference<>() {
-		};
-
-		CompletableFuture<McpSchema.JSONRPCResponse> mockResponseFuture = new CompletableFuture<>();
-		mockResponseFuture.complete(new JSONRPCResponse("2.0", "1", "result", null));
-		// when(mockTransport.sendMessage(any())).then(invocation -> null);
-		when(mockObjectMapper.convertValue("result", typeRef)).thenReturn("result");
-
-		CompletableFuture<String> result = mcpSession.sendRequest(method, params, typeRef);
-		assertEquals("result", result.get());
-	}
-
-	@Test
-	void sendRequestShouldHandleTimeout() {
-		String method = "timeoutMethod";
-		Object params = Map.of("key", "value");
-		TypeReference<String> typeRef = new TypeReference<>() {
-		};
-
-		CompletableFuture<String> result = mcpSession.sendRequest(method, params, typeRef);
-		assertThrows(Exception.class, result::get);
-	}
-
-	@Test
-	void closeShouldInvokeTransportClose() {
-		mcpSession.close();
-		verify(mockTransport, times(1)).close();
-	}
-
-}
diff --git a/mcp/src/test/java/spring/ai/mcp/spec/McpTransportTest.java b/mcp/src/test/java/spring/ai/mcp/spec/McpTransportTest.java
deleted file mode 100644
index 07be6fd..0000000
--- a/mcp/src/test/java/spring/ai/mcp/spec/McpTransportTest.java
+++ /dev/null
@@ -1,128 +0,0 @@
-/* 
-* 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 org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.mockito.ArgumentCaptor;
-import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
-import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
-import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
-import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
-
-import java.time.Duration;
-import java.util.Map;
-import java.util.function.Consumer;
-
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.Mockito.*;
-
-class McpTransportTest {
-
-    private McpTransport mockTransport;
-    private McpSession mcpSession;
-
-    @BeforeEach
-    void setUp() {
-        mockTransport = mock(McpTransport.class);
-        mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), new ObjectMapper());
-    }
-
-    @Test
-    void constructorShouldInitializeTransport() {
-        // Verify transport setup during construction
-        verify(mockTransport).setMessageHandler(any(Consumer.class));
-        verify(mockTransport).setErrorHandler(any(Consumer.class));
-        verify(mockTransport).start();
-    }
-
-    @Test
-    void sendRequestShouldCallTransportSendMessage() {
-        String method = "testMethod";
-        Object params = Map.of("key", "value");
-
-        // Call sendRequest
-        mcpSession.sendRequest(method, params, new TypeReference<String>() {});
-
-        // Capture the message sent to transport
-        ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
-        verify(mockTransport).sendMessage(messageCaptor.capture());
-
-        // Verify the message details
-        JSONRPCRequest request = (JSONRPCRequest) messageCaptor.getValue();
-        assertEquals("testMethod", request.method());
-        assertEquals("value", ((Map<?, ?>) request.params()).get("key"));
-        assertNotNull(request.id());
-    }
-
-    @Test
-    void sendNotificationShouldCallTransportSendMessage() {
-        String method = "notifyMethod";
-        Map<String, Object> params = Map.of("param1", "value1");
-
-        // Call sendNotification
-        mcpSession.sendNotification(method, params);
-
-        // Capture the notification sent to transport
-        ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
-        verify(mockTransport).sendMessage(messageCaptor.capture());
-
-        // Verify the notification details
-        JSONRPCNotification notification = (JSONRPCNotification) messageCaptor.getValue();
-        assertEquals("notifyMethod", notification.method());
-        assertEquals("value1", ((Map<?, ?>) notification.params()).get("param1"));
-    }
-
-    @Test
-    void messageHandlerShouldProcessResponse() {
-        // Capture the message handler during initialization
-        ArgumentCaptor<Consumer<JSONRPCMessage>> handlerCaptor = ArgumentCaptor.forClass(Consumer.class);
-        verify(mockTransport).setMessageHandler(handlerCaptor.capture());
-        Consumer<JSONRPCMessage> handler = handlerCaptor.getValue();
-
-        // Simulate a JSONRPCResponse
-        JSONRPCResponse response = new JSONRPCResponse("2.0", "1", "result", null);
-        handler.accept(response);
-
-        // Verify the response was processed
-        // (Use CompletableFuture to match the response if needed in real tests)
-    }
-
-    @Test
-    void errorHandlerShouldLogErrors() {
-        // Capture the error handler during initialization
-        ArgumentCaptor<Consumer<String>> handlerCaptor = ArgumentCaptor.forClass(Consumer.class);
-        verify(mockTransport).setErrorHandler(handlerCaptor.capture());
-        Consumer<String> errorHandler = handlerCaptor.getValue();
-
-        // Simulate an error
-        String errorMessage = "Test error";
-        errorHandler.accept(errorMessage);
-
-        // Verify the error was logged or processed
-        // (In real tests, capture logging output or other side effects)
-    }
-
-    @Test
-    void closeShouldStopTransport() {
-        mcpSession.close();
-        verify(mockTransport).close();
-    }
-}
diff --git a/spring/src/main/java/spring/ai/mcp/McpApplication.java b/spring/src/main/java/spring/ai/mcp/McpApplication.java
index 6e79325..10116ae 100644
--- a/spring/src/main/java/spring/ai/mcp/McpApplication.java
+++ b/spring/src/main/java/spring/ai/mcp/McpApplication.java
@@ -6,7 +6,7 @@ import java.util.List;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import spring.ai.mcp.client.McpClient;
 import spring.ai.mcp.client.McpSyncClient;
-import spring.ai.mcp.client.stdio.StdioServerParameters;
+import spring.ai.mcp.client.stdio.ServerParameters;
 import spring.ai.mcp.client.stdio.StdioServerTransport;
 
 import org.springframework.ai.chat.client.ChatClient;
@@ -55,15 +55,14 @@ public class McpApplication {
 	@Bean(destroyMethod = "close")
 	public McpSyncClient clientSession() {
 
-		var stdioParams = StdioServerParameters.builder("npx")
+		var stdioParams = ServerParameters.builder("npx")
 			.args("-y", "@modelcontextprotocol/server-filesystem",
 					"/Users/christiantzolov/Dev/projects/demo/mcp-server/dir")
 			.build();
 
 		McpSyncClient mcpClient = null;
 		try {			
-			mcpClient = McpClient.sync(new StdioServerTransport(stdioParams,
-				Duration.ofMillis(100)),
+			mcpClient = McpClient.sync(new StdioServerTransport(stdioParams),
 					Duration.ofSeconds(10), new ObjectMapper());
 
 			var init = mcpClient.initialize();