feat(core): add roots list support and improve documentation

- Add roots list capability to allow servers to understand filesystem boundaries - Move architecture diagram and add detailed documentation for transports and sessions - Enhance client capabilities with roots and sampling features - Rename withRequestTimeout to requestTimeout for better API consistency - Add tests for roots list functionality Resolves #8
2024-12-19 12:09:10 +01:00 · 2024-12-19 12:09:10 +01:00 · 82c64f5ad3
commit 82c64f5ad3
parent 7a98a935af
19 changed files with 810 additions and 45 deletions
--- a/README.md
+++ b/README.md
@ -6,7 +6,7 @@ Java SDK for the Model Context Protocol (MCP), providing seamless integration be

 Spring AI MCP is an experimental project that provides Java and Spring Framework integration for the [Model Context Protocol](https://modelcontextprotocol.org/docs/concepts/architecture). It enables Java applications to interact with AI models and tools through a standardized interface, supporting both synchronous and asynchronous communication patterns.

-<img src="spring-ai-mcp-architecture.jpg" width="600">
+<img src="docs/spring-ai-mcp-architecture.jpg" width="600">

 ## Modules

--- a/docs/CLIENT.SSE.TRANSPORT.md
+++ b/docs/CLIENT.SSE.TRANSPORT.md
@ -0,0 +1,144 @@
+# SSE Server Transport Implementation for MCP
+
+## Overview
+
+The `SseClientTransport` class implements the Model Context Protocol (MCP) HTTP with SSE transport specification, providing a bidirectional communication channel between clients and servers. This implementation is part of the Spring AI MCP Core library and follows the [MCP HTTP with SSE Transport Specification](https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse).
+
+## Key Features
+
+- **Bidirectional Communication**:
+  - Inbound: Receives messages through SSE connection
+  - Outbound: Sends messages via HTTP POST requests
+
+- **Error Handling**:
+  - Automatic reconnection for transient failures
+  - Graceful shutdown capabilities
+  - Comprehensive error reporting
+
+- **Message Flow**:
+  1. Client establishes SSE connection to `/sse` endpoint
+  2. Server sends 'endpoint' event with message URI
+  3. Client uses provided URI for outbound messages
+
+## Technical Details
+
+### Transport Architecture
+
+The transport consists of two main processing pipelines:
+
+1. **Inbound Processing**:
+   - Establishes SSE connection to `/sse` endpoint
+   - Handles two types of events:
+     - `message`: JSON-RPC protocol messages
+     - `endpoint`: Server-provided URI for outbound messages
+   - Implements automatic retry logic for connection failures
+
+2. **Outbound Processing**:
+   - Sends messages via HTTP POST to server-provided endpoint
+   - Uses dedicated scheduler for sequential message processing
+   - Handles JSON serialization and HTTP communication
+
+### Key Components
+
+- **WebClient**: Handles HTTP communications for both SSE and POST requests
+- **ObjectMapper**: Manages JSON serialization/deserialization
+- **Scheduler**: Ensures sequential outbound message processing
+- **Message Endpoint Sink**: Manages server-provided endpoint URI
+
+## Usage Example
+
+```java
+// Create WebClient builder with base URL
+WebClient.Builder webClientBuilder = WebClient.builder()
+    .baseUrl("http://localhost:3001");
+
+// Initialize transport
+SseClientTransport transport = new SseClientTransport(webClientBuilder);
+
+// Create a JSON-RPC request
+JSONRPCRequest request = new JSONRPCRequest(
+    McpSchema.JSONRPC_VERSION,
+    "method-name",
+    "request-id",
+    Map.of("key", "value")
+);
+
+// Send message
+transport.sendMessage(request)
+    .subscribe(
+        success -> System.out.println("Message sent successfully"),
+        error -> System.err.println("Error sending message: " + error)
+    );
+
+// Graceful shutdown when done
+transport.closeGracefully()
+    .subscribe();
+```
+
+## Implementation Details
+
+### Message Processing
+
+1. **Inbound Messages**:
+   - Received through SSE connection
+   - Automatically deserialized from JSON
+   - Routed to appropriate handlers based on event type
+
+2. **Outbound Messages**:
+   - Serialized to JSON
+   - Sent via HTTP POST
+   - Processed sequentially on dedicated thread
+
+### Error Handling
+
+- **Connection Failures**:
+  - Automatic retry for transient failures
+  - Configurable retry policies
+  - Graceful handling of permanent failures
+
+- **Message Processing Errors**:
+  - Comprehensive error reporting
+  - Clean failure handling
+  - Maintains system stability
+
+### Shutdown Process
+
+The transport implements a graceful shutdown mechanism that:
+1. Marks transport as closing
+2. Disposes of all subscriptions
+3. Cleans up schedulers
+4. Completes message endpoint sink
+
+## Technical Specifications
+
+- **Protocol**: HTTP with Server-Sent Events (SSE)
+- **Message Format**: JSON-RPC 2.0
+- **Content Type**: application/json for outbound, text/event-stream for inbound
+- **Endpoints**:
+  - SSE Connection: `/sse`
+  - Message Endpoint: Dynamically provided by server
+
+## Dependencies
+
+- Spring WebFlux for reactive HTTP client
+- Jackson for JSON processing
+- Project Reactor for reactive streams
+- SLF4J for logging
+
+## Thread Safety
+
+The implementation is thread-safe and designed for concurrent use:
+- Inbound processing uses SSE event loop
+- Outbound processing uses dedicated single-threaded scheduler
+- State management uses thread-safe constructs
+
+## Testing
+
+The implementation includes comprehensive tests covering:
+- Message processing
+- Error handling
+- Connection management
+- Shutdown procedures
+- Edge cases and failure scenarios
+
+For detailed test examples, see `SseClientTransportTests.java`.
--- a/docs/CLIENT.STDIO.TRANSPORT.md
+++ b/docs/CLIENT.STDIO.TRANSPORT.md
@ -0,0 +1,123 @@
+# StdioClientTransport
+
+The `StdioClientTransport` class implements the [MCP Stdio Transport](https://spec.modelcontextprotocol.io/specification/basic/transports/#stdio) specification, providing communication with MCP servers over standard input/output streams.
+
+## Overview
+
+The Stdio transport uses standard input/output streams for bidirectional communication between the client and server:
+- Messages are sent to the server via stdout
+- Messages are received from the server via stdin
+- Error and debug information is sent via stderr
+
+## Message Format
+
+Messages are exchanged as newline-delimited JSON-RPC messages:
+- Each message is a single line of JSON text
+- Messages are terminated with a newline character (`\n`)
+- Messages follow the JSON-RPC 2.0 specification
+- UTF-8 encoding is used for all messages
+
+### Message Types
+
+The transport handles three types of JSON-RPC messages:
+
+1. **Requests** - Messages with both `method` and `id` fields:
+```json
+{
+    "jsonrpc": "2.0",
+    "id": 1,
+    "method": "example.method",
+    "params": {}
+}
+```
+
+2. **Notifications** - Messages with `method` but no `id` field:
+```json
+{
+    "jsonrpc": "2.0",
+    "method": "example.notification",
+    "params": {}
+}
+```
+
+3. **Responses** - Messages with either `result` or `error` fields:
+```json
+{
+    "jsonrpc": "2.0",
+    "id": 1,
+    "result": {}
+}
+```
+
+## Implementation Details
+
+### Process Management
+
+The transport manages a server process with:
+- Command and arguments specified via `ServerParameters`
+- Environment variables for configuration
+- Proper process cleanup on shutdown
+
+### Threading Model
+
+The implementation uses three dedicated threads:
+1. **Inbound** - Reads and processes messages from stdin
+2. **Outbound** - Writes messages to stdout
+3. **Error** - Handles stderr messages
+
+### Error Handling
+
+Errors are handled at multiple levels:
+- Process startup failures
+- Stream I/O errors
+- Message parsing errors
+- Process termination handling
+
+## Usage Example
+
+```java
+// Create server parameters
+ServerParameters params = ServerParameters.builder()
+    .command("path/to/server")
+    .args(List.of("--arg1", "value1"))
+    .env(Map.of("ENV_VAR", "value"))
+    .build();
+
+// Create transport
+StdioClientTransport transport = new StdioClientTransport(params);
+
+// Start transport
+transport.start();
+
+// Send message
+JSONRPCRequest request = new JSONRPCRequest();
+request.setId(1);
+request.setMethod("example.method");
+request.setParams(new HashMap<>());
+transport.send(request);
+
+// Close transport
+transport.closeGracefully().block();
+```
+
+## Error Messages
+
+Error messages from the server process are:
+- Written to stderr
+- Logged at ERROR level
+- Emitted through the error sink for handling
+- UTF-8 encoded
+
+## Debugging
+
+For debugging purposes:
+1. Enable DEBUG level logging to see detailed message flow
+2. Monitor stderr for server process errors
+3. Check process exit codes for abnormal termination
+
+## Implementation Notes
+
+- The transport ensures thread safety through dedicated schedulers
+- Messages are processed asynchronously using Project Reactor
+- Resources are properly cleaned up on shutdown
+- The implementation is resilient to various failure modes
--- a/docs/MCP-SESSION.md
+++ b/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
--- a/docs/class-diagrams.puml
+++ b/docs/class-diagrams.puml
@ -0,0 +1,132 @@
+@startuml Core Components
+
+interface McpTransport {
+  +Mono<Void> connect(Function<Mono<JSONRPCMessage>, 
+                       Mono<JSONRPCMessage>> handler)
+  +Mono<Void> sendMessage()
+  +void close()
+  +Mono<Void> closeGracefully()
+}
+
+interface McpSession {
+  +<T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef)
+  +Mono<Void> sendNotification(String method, Map<String, Object> params)
+  +Mono<Void> closeGracefully()
+  +close()
+}
+
+class DefaultMcpSession {
+  -McpTransport transport
+  -ObjectMapper objectMapper
+}
+
+class McpClient {
+  +{static} async()
+  +{static} sync()
+}
+
+class McpAsyncClient {
+  +Mono<InitializeResult> initialize()
+  +Mono<CallToolResult> callTool()
+  +Mono<ListToolsResult> listTools()
+  +Mono<ListResourcesResult> listResources()
+  +Mono<ReadResourceResult> readResource()
+  +Mono<ListPromptsResult> listPrompts()
+  +Mono<GetPromptResult> getPrompt()
+}
+
+class McpSyncClient {
+  +InitializeResult initialize()
+  +CallToolResult callTool()
+  +ListToolsResult listTools()
+  +ListResourcesResult listResources()
+  +ReadResourceResult readResource()
+  +ListPromptsResult listPrompts()
+  +GetPromptResult getPrompt()
+}
+
+class StdioClientTransport {
+  -Process serverProcess
+  -ServerParameters parameters
+}
+
+class SseClientTransport {
+  -WebClient webClient
+}
+
+class McpSchema {
+  +class ErrorCodes
+  +interface Request
+  +interface JSONRPCMessage
+  +interface ResourceContents
+  +interface Content
+}
+
+class McpError {
+  -int code
+  -String message
+}
+
+McpTransport <|-- StdioClientTransport
+McpTransport <|-- SseClientTransport
+McpSession <|-- DefaultMcpSession
+DefaultMcpSession --o McpAsyncClient
+McpClient ..> McpAsyncClient : creates
+McpClient ..> McpSyncClient : creates
+McpSyncClient o-- McpAsyncClient
+DefaultMcpSession o-- McpTransport
+McpSchema <.. McpSession : uses
+McpError ..> McpSession : throws
+
+@enduml
+
+@startuml Message Flow
+
+package "MCP Schema" {
+  interface JSONRPCMessage {
+    +String jsonrpc()
+  }
+  
+  interface Request {
+  }
+  
+  class InitializeRequest
+  class CallToolRequest
+  class CreateMessageRequest
+  class CompleteRequest
+  class GetPromptRequest
+}
+
+package "Resource Types" {
+  interface ResourceContents {
+    +String uri()
+    +String mimeType()
+  }
+  
+  class TextResourceContents
+  class BlobResourceContents
+  
+  interface Content {
+    +String type()
+  }
+  
+  class TextContent
+  class ImageContent
+  class EmbeddedResource
+}
+
+JSONRPCMessage <|.. Request
+Request <|.. InitializeRequest
+Request <|.. CallToolRequest
+Request <|.. CreateMessageRequest
+Request <|.. CompleteRequest
+Request <|.. GetPromptRequest
+
+ResourceContents <|.. TextResourceContents
+ResourceContents <|.. BlobResourceContents
+
+Content <|.. TextContent
+Content <|.. ImageContent
+Content <|.. EmbeddedResource
+
+@enduml
--- a/docs/spring-ai-mcp-architecture.jpg
+++ b/docs/spring-ai-mcp-architecture.jpg
--- a/docs/spring-ai-mcp-uml-classdiagram.svg
+++ b/docs/spring-ai-mcp-uml-classdiagram.svg
--- a/spring-ai-mcp-core/README.md
+++ b/spring-ai-mcp-core/README.md
@ -15,6 +15,7 @@ This SDK implements the Model Context Protocol, enabling seamless integration wi
  - Tool discovery and execution
  - Resource management with URI templates
  - Resource subscription system
+  - Roots list management and notifications
  - Prompt handling and management
  - Server initialization and ping
 - Multiple transport implementations:
@ -72,6 +73,43 @@ The SSE transport provides:
 - Graceful shutdown handling
 - Configurable JSON serialization

+### Roots List Support
+
+The SDK supports the MCP roots list capability, which allows servers to understand which directories and files they have access to. Clients can provide a list of root directories/files and notify servers when this list changes.
+
+#### Features
+- Define root providers that supply filesystem access boundaries
+- Support for roots list changed notifications
+- Automatic roots list request handling
+
+#### Example with Roots List Configuration
+
+```java
+// Create root providers
+List<Supplier<List<Root>>> rootProviders = List.of(
+    () -> List.of(new Root("file:///workspace/project", "Project Root")),
+    () -> List.of(new Root("file:///workspace/docs", "Documentation"))
+);
+
+// Create async client with roots list support
+McpAsyncClient client = new McpAsyncClient(
+    new StdioClientTransport(params),
+    Duration.ofSeconds(30),
+    rootProviders,  // Configure root providers
+    true           // Enable roots list changed notifications
+);
+
+// Initialize connection
+client.initialize()
+    .doOnSuccess(result -> {
+        // Connection initialized with roots list capability
+        
+        // Notify server when roots list changes
+        return client.sendRootsListChanged();
+    })
+    .subscribe();
+```
+
 ### Sync Client Example

 ```java
@ -171,7 +209,7 @@ The SDK follows a layered architecture with clear separation of concerns:
 - **StdioClientTransport**: Standard I/O based client to server communication
 - **SseClientTransport**: HTTP-based transport using Server-Sent Events for bidirectional client-sever communication

-<img src="docs/spring-ai-mcp-uml-classdiagram.svg" width="600"/>
+<img src="../docs/spring-ai-mcp-uml-classdiagram.svg" width="600"/>

 ### Key Interactions

--- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java
+++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java
@ -16,19 +16,27 @@
 package org.springframework.ai.mcp.client;

 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Supplier;

 import com.fasterxml.jackson.core.type.TypeReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;

 import org.springframework.ai.mcp.spec.DefaultMcpSession;
+import org.springframework.ai.mcp.spec.DefaultMcpSession.RequestHandler;
 import org.springframework.ai.mcp.spec.McpError;
 import org.springframework.ai.mcp.spec.McpSchema;
 import org.springframework.ai.mcp.spec.McpSchema.GetPromptRequest;
 import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult;
 import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult;
 import org.springframework.ai.mcp.spec.McpSchema.PaginatedRequest;
+import org.springframework.ai.mcp.spec.McpSchema.Root;
 import org.springframework.ai.mcp.spec.McpTransport;

 /**
@ -52,23 +60,62 @@ public class McpAsyncClient {
 	private final DefaultMcpSession mcpSession;

 	/**
-	 * Create a new McpAsyncClient with the given transport.
-	 * @param transport the transport to use.
+	 * Roots define the boundaries of where servers can operate within the filesystem,
+	 * allowing them to understand which directories and files they have access to.
+	 * Servers can request the list of roots from supporting clients and receive
+	 * notifications when that list changes.
 	 */
-	public McpAsyncClient(McpTransport transport) {
-		this(transport, Duration.ofSeconds(20));
-	}
+	private McpSchema.ClientCapabilities.RootCapabilities rootCapabilities;

 	/**
 	 * Create a new McpAsyncClient with the given transport and session request-response
 	 * timeout.
 	 * @param transport the transport to use.
 	 * @param requestTimeout the session request-response timeout.
+	 * @param rootsListProviders the list of suppliers that provide the list of roots
+	 * backing the roots list request.
+	 * @param rootsListChangedNotification whether the client supports roots/list_changed
+	 * notification.
 	 */
-	public McpAsyncClient(McpTransport transport, Duration requestTimeout) {
-		this.mcpSession = new DefaultMcpSession(requestTimeout, transport);
+	public McpAsyncClient(McpTransport transport, Duration requestTimeout,
+			List<Supplier<List<Root>>> rootsListProviders, boolean rootsListChangedNotification) {
+
+		Map<String, RequestHandler> requestHanlers = new HashMap<>();
+
+		if (rootsListProviders != null && !rootsListProviders.isEmpty()) {
+			requestHanlers.put("roots/list", rootsListRequestHandler(rootsListProviders));
+			this.rootCapabilities = new McpSchema.ClientCapabilities.RootCapabilities(rootsListChangedNotification);
+		}
+
+		this.mcpSession = new DefaultMcpSession(requestTimeout, transport, requestHanlers, Map.of());
+
 	}

+	private RequestHandler rootsListRequestHandler(List<Supplier<List<Root>>> rootsListProviders) {
+		return new RequestHandler() {
+			@Override
+			public Mono<Object> handle(Object params) {
+				if (rootsListProviders == null || rootsListProviders.isEmpty()) {
+					return Mono.just(new ArrayList<Root>());
+				}
+
+				List<Mono<List<Root>>> monos = rootsListProviders.stream()
+					.map(supplier -> Mono.fromSupplier(supplier).subscribeOn(Schedulers.boundedElastic()))
+					.toList();
+
+				return Mono.zip(monos, arrays -> {
+					List<Root> combinedList = new ArrayList<>();
+					for (Object array : arrays) {
+						@SuppressWarnings("unchecked")
+						List<Root> roots = (List<Root>) array;
+						combinedList.addAll(roots);
+					}
+					return combinedList;
+				});
+			}
+		};
+	};
+
 	/**
 	 * The initialization phase MUST be the first interaction between client and server.
 	 * During this phase, the client and server:
@ -100,8 +147,8 @@ public class McpAsyncClient {
 	public Mono<McpSchema.InitializeResult> initialize() {
 		McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
 				McpSchema.LATEST_PROTOCOL_VERSION,
-				new McpSchema.ClientCapabilities(null, new McpSchema.ClientCapabilities.RootCapabilities(true), null),
-				new McpSchema.Implementation("mcp-java-client", "0.0.1")); // @formatter:on
+				new McpSchema.ClientCapabilities(null, rootCapabilities, null),
+				new McpSchema.Implementation("mcp-java-client", "0.2.0")); // @formatter:on

 		Mono<McpSchema.InitializeResult> result = this.mcpSession.sendRequest("initialize", initializeRequest,
 				new TypeReference<McpSchema.InitializeResult>() {
--- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java
+++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java
@ -17,7 +17,11 @@
 package org.springframework.ai.mcp.client;

 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Supplier;

+import org.springframework.ai.mcp.spec.McpSchema.Root;
 import org.springframework.ai.mcp.spec.McpTransport;
 import org.springframework.ai.mcp.util.Assert;

@ -38,10 +42,12 @@ import org.springframework.ai.mcp.util.Assert;
 * including request timeout and JSON object mapping.
 *
 * <p>
- * Use the builder pattern for flexible client configuration: <pre>{@code
+ * Use the builder pattern for flexible client configuration:
+ *
+ * <pre>{@code
 * McpClient.using(transport)
- *     .withRequestTimeout(Duration.ofSeconds(5))
- *     .sync(); // or .async()
+ * 		.withRequestTimeout(Duration.ofSeconds(5))
+ * 		.sync(); // or .async()
 * }</pre>
 *
 * @author Christian Tzolov
@ -74,6 +80,10 @@ public class McpClient {

 		private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout

+		private boolean rootsListChangedNotification = false;
+
+		private List<Supplier<List<Root>>> rootsListProviders = new ArrayList<>();
+
 		private Builder(McpTransport transport) {
 			Assert.notNull(transport, "Transport must not be null");
 			this.transport = transport;
@ -84,12 +94,22 @@ public class McpClient {
 		 * @param requestTimeout The duration to wait before timing out requests
 		 * @return This builder instance
 		 */
-		public Builder withRequestTimeout(Duration requestTimeout) {
+		public Builder requestTimeout(Duration requestTimeout) {
 			Assert.notNull(requestTimeout, "Request timeout must not be null");
 			this.requestTimeout = requestTimeout;
 			return this;
 		}

+		public Builder rootsListChangedNotification(boolean rootsListChangedNotification) {
+			this.rootsListChangedNotification = rootsListChangedNotification;
+			return this;
+		}
+
+		public Builder rootsListProvider(Supplier<List<Root>> rootsListProvider) {
+			this.rootsListProviders.add(rootsListProvider);
+			return this;
+		}
+
 		/**
 		 * Build a synchronous MCP client.
 		 * @return A new instance of {@link McpSyncClient}
@ -103,7 +123,7 @@ public class McpClient {
 		 * @return A new instance of {@link McpAsyncClient}
 		 */
 		public McpAsyncClient async() {
-			return new McpAsyncClient(transport, requestTimeout);
+			return new McpAsyncClient(transport, requestTimeout, rootsListProviders, rootsListChangedNotification);
 		}

 	}
@ -121,7 +141,7 @@ public class McpClient {
 	 */
 	@Deprecated
 	public static McpAsyncClient async(McpTransport transport, Duration requestTimeout) {
-		return using(transport).withRequestTimeout(requestTimeout).async();
+		return using(transport).requestTimeout(requestTimeout).async();
 	}

 	/**
@ -137,7 +157,7 @@ public class McpClient {
 	 */
 	@Deprecated
 	public static McpSyncClient sync(McpTransport transport, Duration requestTimeout) {
-		return using(transport).withRequestTimeout(requestTimeout).sync();
+		return using(transport).requestTimeout(requestTimeout).sync();
 	}

 }
--- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java
+++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java
@ -108,15 +108,6 @@ public class DefaultMcpSession implements McpSession {

 	}

-	/**
-	 * Creates a new DefaultMcpSession with the specified configuration.
-	 * @param requestTimeout Duration to wait for responses
-	 * @param transport Transport implementation for message exchange
-	 */
-	public DefaultMcpSession(Duration requestTimeout, McpTransport transport) {
-		this(requestTimeout, transport, Map.of(), Map.of());
-	}
-
 	/**
 	 * Creates a new DefaultMcpSession with the specified configuration and handlers.
 	 * @param requestTimeout Duration to wait for responses
--- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java
+++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java
@ -155,16 +155,53 @@ public class McpSchema {
 		@JsonProperty("instructions") String instructions) {
 	} // @formatter:on

+	/**
+	 * Clients can implement additional features to enrich connected MCP servers with
+	 * additional capabilities. These capabilities can be used to extend the functionality
+	 * of the server, or to provide additional information to the server about the
+	 * client's capabilities.
+	 *
+	 * @param experimental WIP
+	 * @param roots define the boundaries of where servers can operate within the
+	 * filesystem, allowing them to understand which directories and files they have
+	 * access to.
+	 * @param sampling Provides a standardized way for servers to request LLM sampling
+	 * (“completions” or “generations”) from language models via clients.
+	 *
+	 */
 	@JsonInclude(JsonInclude.Include.NON_ABSENT)
 	public record ClientCapabilities( // @formatter:off
 		@JsonProperty("experimental") Map<String, Object> experimental,
 		@JsonProperty("roots") RootCapabilities roots,
-		@JsonProperty("sampling") Object sampling) {
+		@JsonProperty("sampling") Sampling sampling) {

+		/**
+		 * Roots define the boundaries of where servers can operate within the filesystem,
+		 * allowing them to understand which directories and files they have access to.
+		 * Servers can request the list of roots from supporting clients and
+		 * receive notifications when that list changes.
+		 *
+		 * @param listChanged Whether the client would send notification about roots
+		 * 		  has changed since the last time the server checked.
+		 */
 		@JsonInclude(JsonInclude.Include.NON_ABSENT)			
 		public record RootCapabilities(
 			@JsonProperty("listChanged") Boolean listChanged) {
 		}
+
+		/**
+		 * Provides a standardized way for servers to request LLM
+	 	 * sampling (“completions” or “generations”) from language
+		 * models via clients. This flow allows clients to maintain
+		 * control over model access, selection, and permissions
+		 * while enabling servers to leverage AI capabilities—with
+		 * no server API keys necessary. Servers can request text or
+		 * image-based interactions and optionally include context
+		 * from MCP servers in their prompts.
+		 */
+		@JsonInclude(JsonInclude.Include.NON_ABSENT)			
+		public record Sampling() {
+		}
 	}// @formatter:on

 	@JsonInclude(JsonInclude.Include.NON_ABSENT)
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests2.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests2.java
@ -39,7 +39,7 @@ public class ClientSessionTests2 {
 			.build();

 		try (McpSyncClient clientSession = McpClient.using(new StdioClientTransport(stdioParams))
-			.withRequestTimeout(Duration.ofSeconds(10))
+			.requestTimeout(Duration.ofSeconds(10))
 			.sync()) {

 			clientSession.initialize();
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests3.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests3.java
@ -42,7 +42,7 @@ public class ClientSessionTests3 {
 		try {

 			clientSession = McpClient.using(new StdioClientTransport(stdioParams))
-				.withRequestTimeout(Duration.ofSeconds(10))
+				.requestTimeout(Duration.ofSeconds(10))
 				.sync();

 			clientSession.initialize();
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/MainSSE.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/MainSSE.java
@ -40,7 +40,7 @@ public class MainSSE {

 		try (McpSyncClient mcpClient = McpClient
 			.using(new SseClientTransport(WebClient.builder().baseUrl("http://localhost:3001")))
-			.withRequestTimeout(Duration.ofSeconds(1000))
+			.requestTimeout(Duration.ofSeconds(1000))
 			.sync()) {

 			mcpClient.initialize();
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java
@ -17,7 +17,9 @@
 package org.springframework.ai.mcp.client;

 import java.time.Duration;
+import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;

 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@ -28,6 +30,7 @@ import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
 import org.springframework.ai.mcp.spec.McpSchema.GetPromptRequest;
 import org.springframework.ai.mcp.spec.McpSchema.Prompt;
 import org.springframework.ai.mcp.spec.McpSchema.Resource;
+import org.springframework.ai.mcp.spec.McpSchema.Root;
 import org.springframework.ai.mcp.spec.McpSchema.Tool;
 import org.springframework.ai.mcp.spec.McpTransport;

@ -64,7 +67,7 @@ public abstract class AbstractMcpAsyncClientTests {
 		this.mcpTransport = createMcpTransport();

 		assertThatCode(() -> {
-			mcpAsyncClient = McpClient.using(mcpTransport).withRequestTimeout(TIMEOUT).async();
+			mcpAsyncClient = McpClient.using(mcpTransport).requestTimeout(TIMEOUT).async();
 			mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
 		}).doesNotThrowAnyException();
 	}
@ -83,7 +86,7 @@ public abstract class AbstractMcpAsyncClientTests {
 		assertThatThrownBy(() -> McpClient.using(null).sync()).isInstanceOf(IllegalArgumentException.class)
 			.hasMessage("Transport must not be null");

-		assertThatThrownBy(() -> McpClient.using(mcpTransport).withRequestTimeout(null).sync())
+		assertThatThrownBy(() -> McpClient.using(mcpTransport).requestTimeout(null).sync())
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessage("Request timeout must not be null");
 	}
@ -123,11 +126,6 @@ public abstract class AbstractMcpAsyncClientTests {
 		assertThatThrownBy(() -> mcpAsyncClient.callTool(invalidRequest).block()).isInstanceOf(Exception.class);
 	}

-	@Test
-	void testRootsListChanged() {
-		assertThatCode(() -> mcpAsyncClient.sendRootsListChanged().block()).doesNotThrowAnyException();
-	}
-
 	@Test
 	void testListResources() {
 		StepVerifier.create(mcpAsyncClient.listResources(null)).consumeNextWith(resources -> {
@ -175,4 +173,21 @@ public abstract class AbstractMcpAsyncClientTests {
 			.verifyComplete();
 	}

+	@Test
+	void testRootsListChanged() {
+		assertThatCode(() -> mcpAsyncClient.sendRootsListChanged().block()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testInitializeWithRootsListProviders() {
+		var transport = createMcpTransport();
+		List<Supplier<List<Root>>> providers = List.of(() -> List.of(new Root("file:///test/path", "test-root")));
+
+		var client = new McpAsyncClient(transport, TIMEOUT, providers, true);
+
+		assertThatCode(() -> client.initialize().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+
+		assertThatCode(() -> client.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
 }
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java
@ -64,7 +64,7 @@ public abstract class AbstractMcpSyncClientTests {
 		this.mcpTransport = createMcpTransport();

 		assertThatCode(() -> {
-			mcpSyncClient = McpClient.using(mcpTransport).withRequestTimeout(TIMEOUT).sync();
+			mcpSyncClient = McpClient.using(mcpTransport).requestTimeout(TIMEOUT).sync();
 			mcpSyncClient.initialize();
 		}).doesNotThrowAnyException();
 	}
@ -82,7 +82,7 @@ public abstract class AbstractMcpSyncClientTests {
 		assertThatThrownBy(() -> McpClient.using(null).sync()).isInstanceOf(IllegalArgumentException.class)
 			.hasMessage("Transport must not be null");

-		assertThatThrownBy(() -> McpClient.using(mcpTransport).withRequestTimeout(null).sync())
+		assertThatThrownBy(() -> McpClient.using(mcpTransport).requestTimeout(null).sync())
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessage("Request timeout must not be null");
 	}
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java
@ -0,0 +1,140 @@
+/*
+ * 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.client;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+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 reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
+import reactor.core.scheduler.Schedulers;
+
+import org.springframework.ai.mcp.spec.McpSchema;
+import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification;
+import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
+import org.springframework.ai.mcp.spec.McpSchema.Root;
+import org.springframework.ai.mcp.spec.McpTransport;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class McpAsyncClientResponseHandlerTests {
+
+	private McpAsyncClient asyncMcpClient;
+
+	private MockMcpTransport transport;
+
+	@SuppressWarnings("unused")
+	private static class MockMcpTransport implements McpTransport {
+
+		private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
+
+		private Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
+
+		private Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
+
+		private Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
+
+		public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
+			if (inbound.tryEmitNext(message).isFailure()) {
+				throw new RuntimeException("Failed to emit message " + message);
+			}
+			inboundMessageCount.incrementAndGet();
+		}
+
+		@Override
+		public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
+			if (outgoing.tryEmitNext(message).isFailure()) {
+				return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
+			}
+			return Mono.empty();
+		}
+
+		public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
+			return (JSONRPCRequest) outboundView.blockFirst();
+		}
+
+		public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
+			return (JSONRPCNotification) outboundView.blockFirst();
+		}
+
+		public McpSchema.JSONRPCMessage getLastSentMessage() {
+			return outboundView.blockFirst();
+		}
+
+		@Override
+		public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
+			return inbound.asFlux()
+				.publishOn(Schedulers.boundedElastic())
+				.flatMap(message -> Mono.just(message).transform(handler))
+				.then();
+		}
+
+		@Override
+		public Mono<Void> closeGracefully() {
+			return Mono.empty();
+		}
+
+		@Override
+		public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
+			return new ObjectMapper().convertValue(data, typeRef);
+		}
+
+	}
+
+	@BeforeEach
+	void setUp() {
+		transport = new MockMcpTransport();
+
+		Supplier<List<Root>> rootsListProvider = () -> List.of(new Root("file:///test/path", "test-root"));
+
+		asyncMcpClient = McpClient.using(transport).rootsListProvider(rootsListProvider).async();
+	}
+
+	@AfterEach
+	void tearDown() {
+		if (asyncMcpClient != null) {
+			asyncMcpClient.closeGracefully();
+		}
+	}
+
+	@Test
+	void testRootsListRequestHandling() {
+
+		// Simulate incoming request
+		McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "roots/list",
+				"test-id", null);
+		transport.simulateIncomingMessage(request);
+
+		// Verify response
+		McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
+		assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
+
+		McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
+		assertThat(response.id()).isEqualTo("test-id");
+		assertThat(response.result()).isEqualTo(List.of(new Root("file:///test/path", "test-root")));
+		assertThat(response.error()).isNull();
+	}
+
+}
--- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java
+++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java
@ -62,8 +62,6 @@ class DefaultMcpSessionTests {

 	private MockMcpTransport transport;

-	private ObjectMapper objectMapper;
-
 	@SuppressWarnings("unused")
 	private static class MockMcpTransport implements McpTransport {

@ -124,7 +122,6 @@ class DefaultMcpSessionTests {

 	@BeforeEach
 	void setUp() {
-		objectMapper = new ObjectMapper();
 		transport = new MockMcpTransport();
 		session = new DefaultMcpSession(TIMEOUT, transport, Map.of(),
 				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params))));
@ -139,10 +136,12 @@ class DefaultMcpSessionTests {

 	@Test
 	void testConstructorWithInvalidArguments() {
-		assertThatThrownBy(() -> new DefaultMcpSession(null, transport)).isInstanceOf(IllegalArgumentException.class)
+		assertThatThrownBy(() -> new DefaultMcpSession(null, transport, Map.of(), Map.of()))
+			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("requstTimeout can not be null");

-		assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, null)).isInstanceOf(IllegalArgumentException.class)
+		assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, null, Map.of(), Map.of()))
+			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("transport can not be null");
 	}