refactor: improve MCP client timeout handling and test structure

Add separate initializationTimeout configuration parameter distinct from requestTimeout
Change ping() method to return Void instead of Object for better API clarity
Refactor tests to consistently use StepVerifier instead of mixing assertion styles
Rename ServletSse* test classes to HttpSse* for better naming consistency
Remove redundant getTimeoutDuration() method in WebFlux tests

Signed-off-by: Christian Tzolov <[email protected]>

Author: tzolov

mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -88,7 +88,6 @@ public class McpAsyncClient {
 
 	/**
 	 * The max timeout to await for the client-server connection to be initialized.
-	 * Usually x2 the request timeout. // TODO should we make it configurable?
 	 */
 	private final Duration initializationTimeout;
 
@@ -151,18 +150,21 @@ public class McpAsyncClient {
 	 * timeout.
 	 * @param transport the transport to use.
 	 * @param requestTimeout the session request-response timeout.
+	 * @param initializationTimeout the max timeout to await for the client-server
 	 * @param features the MCP Client supported features.
 	 */
-	McpAsyncClient(ClientMcpTransport transport, Duration requestTimeout, McpClientFeatures.Async features) {
+	McpAsyncClient(ClientMcpTransport transport, Duration requestTimeout, Duration initializationTimeout,
+			McpClientFeatures.Async features) {
 
 		Assert.notNull(transport, "Transport must not be null");
 		Assert.notNull(requestTimeout, "Request timeout must not be null");
+		Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
 
 		this.clientInfo = features.clientInfo();
 		this.clientCapabilities = features.clientCapabilities();
 		this.transport = transport;
 		this.roots = new ConcurrentHashMap<>(features.roots());
-		this.initializationTimeout = requestTimeout.multipliedBy(2);
+		this.initializationTimeout = initializationTimeout;
 
 		// Request Handlers
 		Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();
@@ -367,12 +369,13 @@ private <T> Mono<T> withInitializationCheck(String actionName,
 
 	/**
 	 * Sends a ping request to the server.
-	 * @return A Mono that completes with the server's ping response
+	 * @return A Mono that completes when the server responds to the ping
 	 */
-	public Mono<Object> ping() {
+	public Mono<Void> ping() {
 		return this.withInitializationCheck("pinging the server", initializedResult -> this.mcpSession
 			.sendRequest(McpSchema.METHOD_PING, null, new TypeReference<Object>() {
-			}));
+			})
+			.then());
 	}
 
 	// --------------------------
@@ -771,7 +774,9 @@ private NotificationHandler asyncLoggingNotificationHandler(
 	 * @see McpSchema.LoggingLevel
 	 */
 	public Mono<Void> setLoggingLevel(LoggingLevel loggingLevel) {
-		Assert.notNull(loggingLevel, "Logging level must not be null");
+		if (loggingLevel == null) {
+			return Mono.error(new McpError("Logging level must not be null"));
+		}
 
 		return this.withInitializationCheck("setting logging level", initializedResult -> {
 			String levelName = this.transport.unmarshalFrom(loggingLevel, new TypeReference<String>() {

mcp/src/main/java/io/modelcontextprotocol/client/McpClient.java

@@ -157,6 +157,8 @@ class SyncSpec {
 
 		private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout
 
+		private Duration initializationTimeout = Duration.ofSeconds(20);
+
 		private ClientCapabilities capabilities;
 
 		private Implementation clientInfo = new Implementation("Java SDK MCP Client", "1.0.0");
@@ -193,6 +195,18 @@ public SyncSpec requestTimeout(Duration requestTimeout) {
 			return this;
 		}
 
+		/**
+		 * @param initializationTimeout The duration to wait for the initializaiton
+		 * lifecycle step to complete.
+		 * @return This builder instance for method chaining
+		 * @throws IllegalArgumentException if initializationTimeout is null
+		 */
+		public SyncSpec initializationTimeout(Duration initializationTimeout) {
+			Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
+			this.initializationTimeout = initializationTimeout;
+			return this;
+		}
+
 		/**
 		 * Sets the client capabilities that will be advertised to the server during
 		 * connection initialization. Capabilities define what features the client
@@ -354,7 +368,8 @@ public McpSyncClient build() {
 
 			McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
 
-			return new McpSyncClient(new McpAsyncClient(transport, this.requestTimeout, asyncFeatures));
+			return new McpSyncClient(
+					new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures));
 		}
 
 	}
@@ -381,6 +396,8 @@ class AsyncSpec {
 
 		private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout
 
+		private Duration initializationTimeout = Duration.ofSeconds(20);
+
 		private ClientCapabilities capabilities;
 
 		private Implementation clientInfo = new Implementation("Spring AI MCP Client", "0.3.1");
@@ -417,6 +434,18 @@ public AsyncSpec requestTimeout(Duration requestTimeout) {
 			return this;
 		}
 
+		/**
+		 * @param initializationTimeout The duration to wait for the initializaiton
+		 * lifecycle step to complete.
+		 * @return This builder instance for method chaining
+		 * @throws IllegalArgumentException if initializationTimeout is null
+		 */
+		public AsyncSpec initializationTimeout(Duration initializationTimeout) {
+			Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
+			this.initializationTimeout = initializationTimeout;
+			return this;
+		}
+
 		/**
 		 * Sets the client capabilities that will be advertised to the server during
 		 * connection initialization. Capabilities define what features the client
@@ -574,7 +603,7 @@ public AsyncSpec loggingConsumers(
 		 * @return a new instance of {@link McpAsyncClient}.
 		 */
 		public McpAsyncClient build() {
-			return new McpAsyncClient(this.transport, this.requestTimeout,
+			return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout,
 					new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
 							this.toolsChangeConsumers, this.resourcesChangeConsumers, this.promptsChangeConsumers,
 							this.loggingConsumers, this.samplingHandler));

mcp/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java

@@ -179,11 +179,10 @@ public void removeRoot(String rootUri) {
 	}
 
 	/**
-	 * Send a synchronous ping request.
-	 * @return
+	 * Send a synchronous ping request to the server.
 	 */
-	public Object ping() {
-		return this.delegate.ping().block();
+	public void ping() {
+		this.delegate.ping().block();
 	}
 
 	// --------------------------