Enhance DefaultMcpSession with request/response handling

- Add comprehensive documentation for MCP session components in MCP-SESSION.md
- Add logging and error handling in McpAsyncClient
- Update test configurations and dependencies
- Add reactor-netty-http and spring-context test dependencies

Resolves #6

pom.xml

@@ -75,7 +75,7 @@
 		<slf4j-api.version>2.0.16</slf4j-api.version>
 		<logback.version>1.5.12</logback.version>
 		<jackson.version>2.18.2</jackson.version>
-		<springframework.version>6.1.13</springframework.version>
+		<springframework.version>6.2.1</springframework.version>
 
 		<!-- plugin versions -->
 		<maven-compiler-plugin.version>3.11.0</maven-compiler-plugin.version>

spring-ai-mcp-core/docs/MCP-SESSION.md

@@ -0,0 +1,78 @@
+# MCP Session Components
+
+## Overview
+
+The MCP (Model Context Protocol) session components implement the [MCP specification](https://spec.modelcontextprotocol.io/) for managing bidirectional JSON-RPC communication between clients and model servers. 
+The implementation provides a robust foundation for request-response patterns and notification handling.
+
+## Core Components
+
+### McpSession Interface
+
+The `McpSession` interface defines the contract for MCP communication with three primary operations:
+
+1. Request-Response Communication
+```java
+<T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef);
+```
+
+2. Notification Handling
+```java
+Mono<Void> sendNotification(String method, Map<String, Object> params);
+```
+
+3. Session Lifecycle Management
+```java
+Mono<Void> closeGracefully();
+```
+
+### DefaultMcpSession Implementation
+
+The `DefaultMcpSession` provides the concrete implementation of the MCP session with the following key features:
+
+#### 1. Message Correlation
+- Uses unique request IDs combining session UUID prefix and atomic counter
+- Maintains thread-safe tracking of pending responses
+- Implements timeout management for requests
+
+#### 2. Handler Registration
+```java
+public interface RequestHandler {
+    Mono<Object> handle(Object params);
+}
+
+public interface NotificationHandler {
+    Mono<Void> handle(Object params);
+}
+```
+
+#### 3. Transport Integration
+- Abstracts transport details through `McpTransport` interface
+- Supports different transport implementations (SSE, STDIO)
+- Manages message serialization using Jackson
+
+## Implementation Details
+
+- The session processes three types of messages: McpSchema.JSONRPCResponse, McpSchema.JSONRPCRequest, McpSchema.JSONRPCNotificaiton.
+- The implementation ensures thread safety through
+
+### Error Handling
+
+WIP
+
+## Integration Guidelines
+
+1. **Transport Selection**
+   - Choose appropriate transport (SSE, STDIO) based on use case
+   - Configure transport-specific parameters
+   - Handle transport lifecycle properly
+
+2. **Message Handling**
+   - Register handlers before starting session
+   - Use appropriate error codes from MCP specification
+   - Implement proper request/response correlation
+
+3. **Type Safety**
+   - Use appropriate TypeReference for response deserialization
+   - Validate request parameters
+   - Handle type conversion errors

spring-ai-mcp-core/pom.xml

@@ -44,12 +44,30 @@
 			<artifactId>reactor-core</artifactId>
 		</dependency>
 
+
 		<dependency>
 			<groupId>org.springframework</groupId>
 			<artifactId>spring-webflux</artifactId>
 			<version>${springframework.version}</version>
 		</dependency>
 
+
+		<dependency>
+			<groupId>io.projectreactor.netty</groupId>
+			<artifactId>reactor-netty-http</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+		<!-- The Spring Context is required due to the reactor-netty connector being dependant on
+		the Spring Lifecycle, as discussed here:
+		https://github.com/spring-projects/spring-framework/issues/31180 -->
+		<dependency>
+			<groupId>org.springframework</groupId>
+			<artifactId>spring-context</artifactId>
+			<version>${springframework.version}</version>
+			<scope>test</scope>
+		</dependency>
+
 		<dependency>
 			<groupId>org.assertj</groupId>
 			<artifactId>assertj-core</artifactId>
@@ -83,4 +101,4 @@
 	</dependencies>
 
 
-</project>
+</project>

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java

@@ -19,6 +19,8 @@ import java.time.Duration;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
 
 import org.springframework.ai.mcp.spec.DefaultMcpSession;
@@ -36,6 +38,8 @@ import org.springframework.ai.mcp.spec.McpTransport;
  */
 public class McpAsyncClient extends DefaultMcpSession {
 
+	private final static Logger logger = LoggerFactory.getLogger(McpAsyncClient.class);
+
 	private static TypeReference<Void> VOID_TYPE_REFERENCE = new TypeReference<>() {
 	};
 
@@ -86,6 +90,11 @@ public class McpAsyncClient extends DefaultMcpSession {
 				});
 
 		return result.flatMap(initializeResult -> {
+
+			logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
+					initializeResult.protocolVersion(), initializeResult.capabilities(), initializeResult.serverInfo(),
+					initializeResult.instructions());
+
 			if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(initializeResult.protocolVersion())) {
 				return Mono.error(new McpError(
 						"Unsupported protocol version from the server: " + initializeResult.protocolVersion()));

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java

@@ -20,6 +20,7 @@ import java.time.Duration;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -31,44 +32,117 @@ import reactor.core.publisher.MonoSink;
 import org.springframework.ai.mcp.client.util.Assert;
 
 /**
- * Implementation of the MCP client session.
+ * Default implementation of the MCP (Model Context Protocol) session that manages
+ * bidirectional JSON-RPC communication between clients and servers. This implementation
+ * follows the MCP specification for message exchange and transport handling.
+ *
+ * <p>
+ * The session manages:
+ * <ul>
+ * <li>Request/response handling with unique message IDs</li>
+ * <li>Notification processing</li>
+ * <li>Message timeout management</li>
+ * <li>Transport layer abstraction</li>
+ * </ul>
  *
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
 public class DefaultMcpSession implements McpSession {
 
+	/** Logger for this class */
 	private static final Logger logger = LoggerFactory.getLogger(DefaultMcpSession.class);
 
-	private final ConcurrentHashMap<Object, MonoSink<McpSchema.JSONRPCResponse>> pendingResponses = new ConcurrentHashMap<>();
-
+	/** Duration to wait for request responses before timing out */
 	private final Duration requestTimeout;
 
+	/** JSON object mapper for serialization/deserialization */
 	private final ObjectMapper objectMapper;
 
+	/** Transport layer implementation for message exchange */
 	private final McpTransport transport;
 
+	/** Map of pending responses keyed by request ID */
+	private final ConcurrentHashMap<Object, MonoSink<McpSchema.JSONRPCResponse>> pendingResponses = new ConcurrentHashMap<>();
+
+	/** Map of request handlers keyed by method name */
 	private final ConcurrentHashMap<String, RequestHandler> requestHandlers = new ConcurrentHashMap<>();
 
+	/** Map of notification handlers keyed by method name */
+	private final ConcurrentHashMap<String, NotificationHandler> notificationHandlers = new ConcurrentHashMap<>();
+
+	/** Session-specific prefix for request IDs */
+	private final String sessionPrefix = UUID.randomUUID().toString().substring(0, 8);
+
+	/** Atomic counter for generating unique request IDs */
+	private final AtomicLong requestCounter = new AtomicLong(0);
+
+	/**
+	 * Functional interface for handling incoming JSON-RPC requests. Implementations
+	 * should process the request parameters and return a response.
+	 */
 	@FunctionalInterface
 	public interface RequestHandler {
 
+		/**
+		 * Handles an incoming request with the given parameters.
+		 * @param params The request parameters
+		 * @return A Mono containing the response object
+		 */
 		Mono<Object> handle(Object params);
 
 	}
 
+	/**
+	 * Functional interface for handling incoming JSON-RPC notifications. Implementations
+	 * should process the notification parameters without returning a response.
+	 */
+	@FunctionalInterface
+	public interface NotificationHandler {
+
+		/**
+		 * Handles an incoming notification with the given parameters.
+		 * @param params The notification parameters
+		 * @return A Mono that completes when the notification is processed
+		 */
+		Mono<Void> handle(Object params);
+
+	}
+
+	/**
+	 * Creates a new DefaultMcpSession with the specified configuration.
+	 * @param requestTimeout Duration to wait for responses
+	 * @param objectMapper JSON object mapper for message serialization
+	 * @param transport Transport implementation for message exchange
+	 */
 	public DefaultMcpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport) {
+		this(requestTimeout, objectMapper, transport, Map.of(), Map.of());
+	}
+
+	/**
+	 * Creates a new DefaultMcpSession with the specified configuration and handlers.
+	 * @param requestTimeout Duration to wait for responses
+	 * @param objectMapper JSON object mapper for message serialization
+	 * @param transport Transport implementation for message exchange
+	 * @param requestHandlers Map of method names to request handlers
+	 * @param notificationHandlers Map of method names to notification handlers
+	 */
+	public DefaultMcpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport,
+			Map<String, RequestHandler> requestHandlers, Map<String, NotificationHandler> notificationHandlers) {
 
 		Assert.notNull(objectMapper, "The ObjectMapper can not be null");
 		Assert.notNull(requestTimeout, "The requstTimeout can not be null");
 		Assert.notNull(transport, "The transport can not be null");
+		Assert.notNull(requestHandlers, "The requestHandlers can not be null");
+		Assert.notNull(notificationHandlers, "The notificationHandlers can not be null");
 
 		this.requestTimeout = requestTimeout;
 		this.objectMapper = objectMapper;
 		this.transport = transport;
+		this.requestHandlers.putAll(requestHandlers);
+		this.notificationHandlers.putAll(notificationHandlers);
 
 		this.transport.setInboundMessageHandler(message -> {
-
 			if (message instanceof McpSchema.JSONRPCResponse response) {
 				var sink = pendingResponses.remove(response.id());
 				if (sink == null) {
@@ -88,7 +162,8 @@ public class DefaultMcpSession implements McpSession {
 						});
 			}
 			else if (message instanceof McpSchema.JSONRPCNotification notification) {
-				logger.info("Notifications not yet supported");
+				handleIncomingNotification(notification).subscribe(null,
+						error -> logger.error("Error handling notification: {}", error.getMessage()));
 			}
 		});
 
@@ -97,9 +172,14 @@ public class DefaultMcpSession implements McpSession {
 		this.transport.start();
 	}
 
+	/**
+	 * Handles an incoming JSON-RPC request by routing it to the appropriate handler.
+	 * @param request The incoming JSON-RPC request
+	 * @return A Mono containing the JSON-RPC response
+	 */
 	private Mono<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCRequest request) {
 		return Mono.defer(() -> {
-			var handler = requestHandlers.get(request.method());
+			var handler = this.requestHandlers.get(request.method());
 			if (handler == null) {
 				return Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
 						new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND,
@@ -110,19 +190,47 @@ public class DefaultMcpSession implements McpSession {
 				.map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null))
 				.onErrorResume(error -> Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
 						null, new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
-								error.getMessage(), null))));
+								error.getMessage(), null)))); // TODO: add error message
+																// through the data field
 		});
 	}
 
-	public void registerRequestHandler(String method, RequestHandler handler) {
-		requestHandlers.put(method, handler);
+	/**
+	 * Handles an incoming JSON-RPC notification by routing it to the appropriate handler.
+	 * @param notification The incoming JSON-RPC notification
+	 * @return A Mono that completes when the notification is processed
+	 */
+	private Mono<Void> handleIncomingNotification(McpSchema.JSONRPCNotification notification) {
+		return Mono.defer(() -> {
+			var handler = notificationHandlers.get(notification.method());
+			if (handler == null) {
+				logger.error("No handler registered for notification method: {}", notification.method());
+				return Mono.empty();
+			}
+			return handler.handle(notification.params());
+		});
 	}
 
+	/**
+	 * Generates a unique request ID in a non-blocking way. Combines a session-specific
+	 * prefix with an atomic counter to ensure uniqueness.
+	 * @return A unique request ID string
+	 */
+	private String generateRequestId() {
+		return this.sessionPrefix + "-" + this.requestCounter.getAndIncrement();
+	}
+
+	/**
+	 * Sends a JSON-RPC request and returns the response.
+	 * @param <T> The expected response type
+	 * @param method The method name to call
+	 * @param requestParams The request parameters
+	 * @param typeRef Type reference for response deserialization
+	 * @return A Mono containing the response
+	 */
 	@Override
 	public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
-		// TODO: UUID API is blocking. Consider non-blocking alternatives to generate
-		// the ID.
-		String requestId = UUID.randomUUID().toString();
+		String requestId = this.generateRequestId();
 
 		return Mono.<McpSchema.JSONRPCResponse>create(sink -> {
 			this.pendingResponses.put(requestId, sink);
@@ -131,25 +239,31 @@ public class DefaultMcpSession implements McpSession {
 			this.transport.sendMessage(jsonrpcRequest)
 				// TODO: It's most efficient to create a dedicated Subscriber here
 				.subscribe(v -> {
-				}, e -> {
+				}, error -> {
 					this.pendingResponses.remove(requestId);
-					sink.error(e);
+					sink.error(error);
 				});
-		}).timeout(this.requestTimeout).handle((jsonRpcResponse, s) -> {
+		}).timeout(this.requestTimeout).handle((jsonRpcResponse, sink) -> {
 			if (jsonRpcResponse.error() != null) {
-				s.error(new McpError(jsonRpcResponse.error()));
+				sink.error(new McpError(jsonRpcResponse.error()));
 			}
 			else {
 				if (typeRef.getType().equals(Void.class)) {
-					s.complete();
+					sink.complete();
 				}
 				else {
-					s.next(this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef));
+					sink.next(this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef));
 				}
 			}
 		});
 	}
 
+	/**
+	 * Sends a JSON-RPC notification.
+	 * @param method The method name for the notification
+	 * @param params The notification parameters
+	 * @return A Mono that completes when the notification is sent
+	 */
 	@Override
 	public Mono<Void> sendNotification(String method, Map<String, Object> params) {
 		McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
@@ -157,11 +271,18 @@ public class DefaultMcpSession implements McpSession {
 		return this.transport.sendMessage(jsonrpcNotification);
 	}
 
+	/**
+	 * Closes the session gracefully, allowing pending operations to complete.
+	 * @return A Mono that completes when the session is closed
+	 */
 	@Override
 	public Mono<Void> closeGracefully() {
 		return transport.closeGracefully();
 	}
 
+	/**
+	 * Closes the session immediately, potentially interrupting pending operations.
+	 */
 	@Override
 	public void close() {
 		transport.close();

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java

@@ -650,11 +650,33 @@ public class McpSchema {
 		}
 	}
 
+	// ---------------------------
 	// Roots
+	// ---------------------------
+	/**
+	 * Represents a root directory or file that the server can operate on.
+	 *
+	 * @param uri The URI identifying the root. This *must* start with file:// for now.
+	 * This restriction may be relaxed in future versions of the protocol to allow other
+	 * URI schemes.
+	 * @param name An optional name for the root. This can be used to provide a
+	 * human-readable identifier for the root, which may be useful for display purposes or
+	 * for referencing the root in other parts of the application.
+	 */
 	@JsonInclude(JsonInclude.Include.NON_ABSENT)
 	public record Root( // @formatter:off
 		@JsonProperty("uri") String uri,
 		@JsonProperty("name") String name) {
 	} // @formatter:on
 
+	/**
+	 * The client's response to a roots/list request from the server. This result contains
+	 * an array of Root objects, each representing a root directory or file that the
+	 * server can operate on.
+	 */
+	@JsonInclude(JsonInclude.Include.NON_ABSENT)
+	public record ListRootsResult( // @formatter:off
+		@JsonProperty("roots") List<Root> roots) {
+	} // @formatter:on
+
 }

spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseServerTransportTests.java

@@ -25,7 +25,6 @@ import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.testcontainers.containers.GenericContainer;
@@ -46,7 +45,6 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  *
  * @author Christian Tzolov
  */
-@Disabled
 @Timeout(15)
 class SseServerTransportTests {
 
@@ -67,9 +65,6 @@ class SseServerTransportTests {
 	// Test class to access protected methods
 	static class TestSseServerTransport extends SseServerTransport {
 
-		// private final Sinks.Many<String> endpointSink =
-		// Sinks.many().replay().latest();
-
 		private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
 
 		public TestSseServerTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
@@ -144,7 +139,6 @@ class SseServerTransportTests {
 			.hasMessageContaining("ObjectMapper must not be null");
 	}
 
-	@Disabled
 	@Test
 	void testMessageProcessing() {
 		// Create a test message
@@ -167,7 +161,6 @@ class SseServerTransportTests {
 		assertThat(transport.getInboundMessageCount()).isEqualTo(1);
 	}
 
-	@Disabled
 	@Test
 	void testResponseMessageProcessing() {
 		// Simulate receiving a response message
@@ -189,7 +182,6 @@ class SseServerTransportTests {
 		assertThat(transport.getInboundMessageCount()).isEqualTo(1);
 	}
 
-	@Disabled
 	@Test
 	void testErrorMessageProcessing() {
 		// Simulate receiving an error message
@@ -214,7 +206,6 @@ class SseServerTransportTests {
 		assertThat(transport.getInboundMessageCount()).isEqualTo(1);
 	}
 
-	@Disabled
 	@Test
 	void testNotificationMessageProcessing() {
 		// Simulate receiving a notification message (no id)
@@ -246,7 +237,6 @@ class SseServerTransportTests {
 		assertThat(transport.getInboundMessageCount()).isEqualTo(0);
 	}
 
-	@Disabled
 	@Test
 	void testRetryBehavior() {
 		// Create a WebClient that simulates connection failures
@@ -261,7 +251,6 @@ class SseServerTransportTests {
 		failingTransport.closeGracefully().block();
 	}
 
-	@Disabled
 	@Test
 	void testMultipleMessageProcessing() {
 		// Simulate receiving multiple messages in sequence
@@ -297,7 +286,6 @@ class SseServerTransportTests {
 		assertThat(transport.getInboundMessageCount()).isEqualTo(2);
 	}
 
-	@Disabled
 	@Test
 	void testMessageOrderPreservation() {
 		// Simulate receiving messages in a specific order

spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java

@@ -0,0 +1,267 @@
+/*
+ * 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 org.springframework.ai.mcp.spec;
+
+import java.time.Duration;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import org.springframework.ai.mcp.client.McpAsyncClient;
+import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification;
+import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Test suite for {@link DefaultMcpSession} that verifies its JSON-RPC message handling,
+ * request-response correlation, and notification processing.
+ *
+ * @author Christian Tzolov
+ */
+class DefaultMcpSessionTests {
+
+	private final static Logger logger = LoggerFactory.getLogger(DefaultMcpSessionTests.class);
+
+	private static final Duration TIMEOUT = Duration.ofSeconds(5);
+
+	private static final String TEST_METHOD = "test.method";
+
+	private static final String TEST_NOTIFICATION = "test.notification";
+
+	private static final String ECHO_METHOD = "echo";
+
+	private DefaultMcpSession session;
+
+	private MockMcpTransport transport;
+
+	private ObjectMapper objectMapper;
+
+	@SuppressWarnings("unused")
+	private static class MockMcpTransport extends AbstractMcpTransport {
+
+		private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
+
+		private final AtomicReference<McpSchema.JSONRPCMessage> lastSentMessage = new AtomicReference<>();
+
+		public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
+			inboundMessageCount.incrementAndGet();
+			this.getInboundSink().tryEmitNext(message);
+		}
+
+		@Override
+		public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
+			lastSentMessage.set(message);
+			return Mono.empty();
+		}
+
+		public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
+			return (JSONRPCRequest) lastSentMessage.get();
+		}
+
+		public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
+			return (JSONRPCNotification) lastSentMessage.get();
+		}
+
+		public McpSchema.JSONRPCMessage getLastSentMessage() {
+			return lastSentMessage.get();
+		}
+
+		@Override
+		public Mono<Void> closeGracefully() {
+			return Mono.empty();
+		}
+
+	}
+
+	@BeforeEach
+	void setUp() {
+		objectMapper = new ObjectMapper();
+		transport = new MockMcpTransport();
+		session = new DefaultMcpSession(TIMEOUT, objectMapper, transport, Map.of(),
+				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params))));
+	}
+
+	@AfterEach
+	void tearDown() {
+		if (session != null) {
+			session.close();
+		}
+	}
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> new DefaultMcpSession(null, objectMapper, transport))
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("requstTimeout can not be null");
+
+		assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, null, transport))
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("ObjectMapper can not be null");
+
+		assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, objectMapper, null))
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("transport can not be null");
+	}
+
+	TypeReference<String> responseType = new TypeReference<>() {
+	};
+
+	@Test
+	void testSendRequest() {
+		String testParam = "test parameter";
+		String responseData = "test response";
+
+		// Create a Mono that will emit the response after the request is sent
+		Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam, responseType)
+			.doOnSubscribe(subscription -> {
+				// Wait a bit to ensure the request is sent and captured
+				Mono.delay(Duration.ofMillis(100)).subscribe(ignored -> {
+					McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
+					transport.simulateIncomingMessage(
+							new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null));
+				});
+			});
+
+		// Verify response handling
+		StepVerifier.create(responseMono).consumeNextWith(response -> {
+			// Verify the request was sent
+			McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest();
+			assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCRequest.class);
+			McpSchema.JSONRPCRequest request = (McpSchema.JSONRPCRequest) sentMessage;
+			assertThat(request.method()).isEqualTo(TEST_METHOD);
+			assertThat(request.params()).isEqualTo(testParam);
+			assertThat(response).isEqualTo(responseData);
+		}).verifyComplete();
+	}
+
+	@Test
+	void testSendRequestWithError() {
+		Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType)
+			.doOnSubscribe(subscription -> {
+				// Wait a bit to ensure the request is sent and captured
+				Mono.delay(Duration.ofMillis(100)).subscribe(ignored -> {
+					McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
+					// Simulate error response
+					McpSchema.JSONRPCResponse.JSONRPCError error = new McpSchema.JSONRPCResponse.JSONRPCError(
+							McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null);
+					transport.simulateIncomingMessage(
+							new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error));
+				});
+			});
+
+		// Verify error handling
+		StepVerifier.create(responseMono).expectError(McpError.class).verify();
+	}
+
+	@Test
+	void testRequestTimeout() {
+
+		Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
+
+		// Verify timeout
+		StepVerifier.create(responseMono)
+			.expectError(java.util.concurrent.TimeoutException.class)
+			.verify(TIMEOUT.plusSeconds(1));
+	}
+
+	@Test
+	void testSendNotification() {
+		Map<String, Object> params = Map.of("key", "value");
+		Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params);
+
+		// Verify notification was sent
+		StepVerifier.create(notificationMono).consumeSubscriptionWith(response -> {
+			McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
+			assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCNotification.class);
+			McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification) sentMessage;
+			assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION);
+			assertThat(notification.params()).isEqualTo(params);
+		}).verifyComplete();
+	}
+
+	@Test
+	void testRequestHandling() {
+		String echoMessage = "Hello MCP!";
+		Map<String, DefaultMcpSession.RequestHandler> requestHandlers = Map.of(ECHO_METHOD,
+				params -> Mono.just(params));
+		session = new DefaultMcpSession(TIMEOUT, objectMapper, transport, requestHandlers, Map.of());
+
+		// Simulate incoming request
+		McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD,
+				"test-id", echoMessage);
+		transport.simulateIncomingMessage(request);
+
+		// Verify response
+		McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
+		assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
+		McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
+		assertThat(response.result()).isEqualTo(echoMessage);
+		assertThat(response.error()).isNull();
+	}
+
+	@Test
+	void testNotificationHandling() {
+		AtomicReference<Object> receivedParams = new AtomicReference<>();
+
+		session = new DefaultMcpSession(TIMEOUT, objectMapper, transport, Map.of(),
+				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.set(params))));
+
+		// Simulate incoming notification from the server
+		Map<String, Object> notificationParams = Map.of("status", "ready");
+
+		McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
+				TEST_NOTIFICATION, notificationParams);
+
+		transport.simulateIncomingMessage(notification);
+
+		// Verify handler was called
+		assertThat(receivedParams.get()).isEqualTo(notificationParams);
+	}
+
+	@Test
+	void testUnknownMethodHandling() {
+		// Simulate incoming request for unknown method
+		McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method",
+				"test-id", null);
+		transport.simulateIncomingMessage(request);
+
+		// Verify error response
+		McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
+		assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
+		McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
+		assertThat(response.error()).isNotNull();
+		assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
+	}
+
+	@Test
+	void testGracefulShutdown() {
+		StepVerifier.create(session.closeGracefully()).verifyComplete();
+	}
+
+}