diff --git a/README.md b/README.md index 5d8c079..2f216c5 100644 --- 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. - + ## Modules diff --git a/docs/CLIENT.SSE.TRANSPORT.md b/docs/CLIENT.SSE.TRANSPORT.md new file mode 100644 index 0000000..e65794f --- /dev/null +++ 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`. diff --git a/docs/CLIENT.STDIO.TRANSPORT.md b/docs/CLIENT.STDIO.TRANSPORT.md new file mode 100644 index 0000000..1c2dc6f --- /dev/null +++ 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 diff --git a/docs/MCP-SESSION.md b/docs/MCP-SESSION.md new file mode 100644 index 0000000..7a48505 --- /dev/null +++ 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 + Mono sendRequest(String method, Object requestParams, TypeReference typeRef); +``` + +2. Notification Handling +```java +Mono sendNotification(String method, Map params); +``` + +3. Session Lifecycle Management +```java +Mono 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 handle(Object params); +} + +public interface NotificationHandler { + Mono 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 diff --git a/docs/class-diagrams.puml b/docs/class-diagrams.puml new file mode 100644 index 0000000..4b2c8cd --- /dev/null +++ b/docs/class-diagrams.puml @@ -0,0 +1,132 @@ +@startuml Core Components + +interface McpTransport { + +Mono connect(Function, + Mono> handler) + +Mono sendMessage() + +void close() + +Mono closeGracefully() +} + +interface McpSession { + + Mono sendRequest(String method, Object requestParams, TypeReference typeRef) + +Mono sendNotification(String method, Map params) + +Mono closeGracefully() + +close() +} + +class DefaultMcpSession { + -McpTransport transport + -ObjectMapper objectMapper +} + +class McpClient { + +{static} async() + +{static} sync() +} + +class McpAsyncClient { + +Mono initialize() + +Mono callTool() + +Mono listTools() + +Mono listResources() + +Mono readResource() + +Mono listPrompts() + +Mono 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 diff --git a/spring-ai-mcp-architecture.jpg b/docs/spring-ai-mcp-architecture.jpg similarity index 100% rename from spring-ai-mcp-architecture.jpg rename to docs/spring-ai-mcp-architecture.jpg diff --git a/docs/spring-ai-mcp-uml-classdiagram.svg b/docs/spring-ai-mcp-uml-classdiagram.svg new file mode 100644 index 0000000..6b05be7 --- /dev/null +++ b/docs/spring-ai-mcp-uml-classdiagram.svg @@ -0,0 +1 @@ +McpTransportMono<Void> connect(Function<Mono<JSONRPCMessage>,Mono<JSONRPCMessage>> handler)Mono<Void> sendMessage()void close()Mono<Void> closeGracefully()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()DefaultMcpSessionMcpTransport transportObjectMapper objectMapperMcpClientasync()sync()McpAsyncClientMono<InitializeResult> initialize()Mono<CallToolResult> callTool()Mono<ListToolsResult> listTools()Mono<ListResourcesResult> listResources()Mono<ReadResourceResult> readResource()Mono<ListPromptsResult> listPrompts()Mono<GetPromptResult> getPrompt()McpSyncClientInitializeResult initialize()CallToolResult callTool()ListToolsResult listTools()ListResourcesResult listResources()ReadResourceResult readResource()ListPromptsResult listPrompts()GetPromptResult getPrompt()StdioClientTransportProcess serverProcessServerParameters parametersSseClientTransportWebClient webClientMcpSchemaclass ErrorCodesinterface Requestinterface JSONRPCMessageinterface ResourceContentsinterface ContentMcpErrorint codeString messagecreatescreatesusesthrows \ No newline at end of file diff --git a/spring-ai-mcp-core/README.md b/spring-ai-mcp-core/README.md index c4fb117..60a092f 100644 --- 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>> 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 - + ### Key Interactions diff --git 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 index 7accb87..f5e55a4 100644 --- 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>> rootsListProviders, boolean rootsListChangedNotification) { + + Map 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>> rootsListProviders) { + return new RequestHandler() { + @Override + public Mono handle(Object params) { + if (rootsListProviders == null || rootsListProviders.isEmpty()) { + return Mono.just(new ArrayList()); + } + + List>> monos = rootsListProviders.stream() + .map(supplier -> Mono.fromSupplier(supplier).subscribeOn(Schedulers.boundedElastic())) + .toList(); + + return Mono.zip(monos, arrays -> { + List combinedList = new ArrayList<>(); + for (Object array : arrays) { + @SuppressWarnings("unchecked") + List roots = (List) 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 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 result = this.mcpSession.sendRequest("initialize", initializeRequest, new TypeReference() { diff --git 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 index a9b968b..847849e 100644 --- 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. * *

- * Use the builder pattern for flexible client configuration: 

{@code
+ * Use the builder pattern for flexible client configuration:
+ *
+ * 
{@code
  * McpClient.using(transport)
- *     .withRequestTimeout(Duration.ofSeconds(5))
- *     .sync(); // or .async()
+ * 		.withRequestTimeout(Duration.ofSeconds(5))
+ * 		.sync(); // or .async()
  * }

  *
  * @author Christian Tzolov
@@ -74,6 +80,10 @@ public class McpClient {
 
 		private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout
 
+		private boolean rootsListChangedNotification = false;
+
+		private List>> 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> 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();
 	}
 
 }
diff --git 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
index cb82f54..71786a6 100644
--- 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
diff --git 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
index 1415d9f..86cf60e 100644
--- 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 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)
diff --git 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
index 12c3d71..9ab8205 100644
--- 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();
diff --git 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
index de1417c..474a961 100644
--- 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();
diff --git 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
index 1e96dad..e7a6f8d 100644
--- 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();
diff --git 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
index 08cf29a..eba0302 100644
--- 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>> 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();
+	}
+
 }
diff --git 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
index 00dc68d..eedeb00 100644
--- 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");
 	}
diff --git 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
new file mode 100644
index 0000000..c72adcd
--- /dev/null
+++ 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 outgoing = Sinks.many().multicast().onBackpressureBuffer();
+
+		private Sinks.Many inbound = Sinks.many().unicast().onBackpressureBuffer();
+
+		private Flux 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 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 connect(Function, Mono> handler) {
+			return inbound.asFlux()
+				.publishOn(Schedulers.boundedElastic())
+				.flatMap(message -> Mono.just(message).transform(handler))
+				.then();
+		}
+
+		@Override
+		public Mono closeGracefully() {
+			return Mono.empty();
+		}
+
+		@Override
+		public  T unmarshalFrom(Object data, TypeReference typeRef) {
+			return new ObjectMapper().convertValue(data, typeRef);
+		}
+
+	}
+
+	@BeforeEach
+	void setUp() {
+		transport = new MockMcpTransport();
+
+		Supplier> 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();
+	}
+
+}
diff --git 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
index 75e82a3..a29e0d2 100644
--- 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");
 	}