feat(client): add client tool output schema validation and caching (#302)

- Add JSON schema validation for tool call results against output schemas
- Implement automatic tool output schema caching during initialization
- Add `enableCallToolSchemaCaching` configuration option to enable/disable
  schema caching
- Add `JsonSchemaValidator` integration to McpClient builder APIs
- Introduce post-initialization hook mechanism for performing operations
  after successful client initialization
- Cache tool output schemas during `listTools` operations when caching
  is enabled
- Validate tool results against cached schemas in `callTool` operations
- Return error CallToolResult when validation fails
- Add test coverage
- Convert validateToolResult from Mono to synchronous method
- Throw IllegalArgumentException on validation errors

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>

Co-authored-by: Anurag Pant <pantanurag555@github>
Co-authored-by: Christian Tzolov <christian.tzolov@broadcom.com>
Co-authored-by: Daniel Garnier-Moiroux <git@garnier.wf>
Co-authored-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java

@@ -11,14 +11,13 @@ import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import io.modelcontextprotocol.spec.McpClientSession;
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
 import io.modelcontextprotocol.util.Assert;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
 import reactor.util.context.ContextView;
@@ -99,21 +98,30 @@ class LifecycleInitializer {
 	 */
 	private final Duration initializationTimeout;
 
+	/**
+	 * Post-initialization hook to perform additional operations after every successful
+	 * initialization.
+	 */
+	private final Function<Initialization, Mono<Void>> postInitializationHook;
+
 	public LifecycleInitializer(McpSchema.ClientCapabilities clientCapabilities, McpSchema.Implementation clientInfo,
 			List<String> protocolVersions, Duration initializationTimeout,
-			Function<ContextView, McpClientSession> sessionSupplier) {
+			Function<ContextView, McpClientSession> sessionSupplier,
+			Function<Initialization, Mono<Void>> postInitializationHook) {
 
 		Assert.notNull(sessionSupplier, "Session supplier must not be null");
 		Assert.notNull(clientCapabilities, "Client capabilities must not be null");
 		Assert.notNull(clientInfo, "Client info must not be null");
 		Assert.notEmpty(protocolVersions, "Protocol versions must not be empty");
 		Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
+		Assert.notNull(postInitializationHook, "Post-initialization hook must not be null");
 
 		this.sessionSupplier = sessionSupplier;
 		this.clientCapabilities = clientCapabilities;
 		this.clientInfo = clientInfo;
 		this.protocolVersions = Collections.unmodifiableList(new ArrayList<>(protocolVersions));
 		this.initializationTimeout = initializationTimeout;
+		this.postInitializationHook = postInitializationHook;
 	}
 
 	/**
@@ -148,10 +156,6 @@ class LifecycleInitializer {
 
 	}
 
-	/**
-	 * Default implementation of the {@link Initialization} interface that manages the MCP
-	 * client initialization process.
-	 */
 	private static class DefaultInitialization implements Initialization {
 
 		/**
@@ -199,29 +203,20 @@ class LifecycleInitializer {
 			this.mcpClientSession.set(mcpClientSession);
 		}
 
-		/**
-		 * Returns a Mono that completes when the MCP client initialization is complete.
-		 * This allows subscribers to wait for the initialization to finish before
-		 * proceeding with further operations.
-		 * @return A Mono that emits the result of the MCP initialization process
-		 */
 		private Mono<McpSchema.InitializeResult> await() {
 			return this.initSink.asMono();
 		}
 
-		/**
-		 * Completes the initialization process with the given result. It caches the
-		 * result and emits it to all subscribers waiting for the initialization to
-		 * complete.
-		 * @param initializeResult The result of the MCP initialization process
-		 */
 		private void complete(McpSchema.InitializeResult initializeResult) {
-			// first ensure the result is cached
-			this.result.set(initializeResult);
 			// inform all the subscribers waiting for the initialization
 			this.initSink.emitValue(initializeResult, Sinks.EmitFailureHandler.FAIL_FAST);
 		}
 
+		private void cacheResult(McpSchema.InitializeResult initializeResult) {
+			// first ensure the result is cached
+			this.result.set(initializeResult);
+		}
+
 		private void error(Throwable t) {
 			this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST);
 		}
@@ -263,7 +258,7 @@ class LifecycleInitializer {
 			}
 			// Providing an empty operation since we are only interested in triggering
 			// the implicit initialization step.
-			withIntitialization("re-initializing", result -> Mono.empty()).subscribe();
+			this.withInitialization("re-initializing", result -> Mono.empty()).subscribe();
 		}
 	}
 
@@ -275,7 +270,7 @@ class LifecycleInitializer {
 	 * @param operation The operation to execute when the client is initialized
 	 * @return A Mono that completes with the result of the operation
 	 */
-	public <T> Mono<T> withIntitialization(String actionName, Function<Initialization, Mono<T>> operation) {
+	public <T> Mono<T> withInitialization(String actionName, Function<Initialization, Mono<T>> operation) {
 		return Mono.deferContextual(ctx -> {
 			DefaultInitialization newInit = new DefaultInitialization();
 			DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit);
@@ -283,8 +278,8 @@ class LifecycleInitializer {
 			boolean needsToInitialize = previous == null;
 			logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
 
-			Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit, ctx)
+			Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
-					: previous.await();
+					? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await();
 
 			return initializationJob.map(initializeResult -> this.initializationRef.get())
 				.timeout(this.initializationTimeout)
@@ -296,7 +291,9 @@ class LifecycleInitializer {
 		});
 	}
 
-	private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization, ContextView ctx) {
+	private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization,
+			Function<Initialization, Mono<Void>> postInitOperation, ContextView ctx) {
+
 		initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
 
 		McpClientSession mcpClientSession = initialization.mcpSession();
@@ -323,6 +320,9 @@ class LifecycleInitializer {
 
 			return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
 				.thenReturn(initializeResult);
+		}).flatMap(initializeResult -> {
+			initialization.cacheResult(initializeResult);
+			return postInitOperation.apply(initialization).thenReturn(initializeResult);
 		}).doOnNext(initialization::complete).onErrorResume(ex -> {
 			initialization.error(ex);
 			return Mono.error(ex);

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

@@ -15,12 +15,12 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
-import org.slf4j.Logger;
+import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
-import org.slf4j.LoggerFactory;
-
 import io.modelcontextprotocol.json.TypeRef;
-
+import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
 import io.modelcontextprotocol.spec.McpClientSession;
+import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler;
+import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
@@ -35,10 +35,10 @@ import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
 import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
 import io.modelcontextprotocol.spec.McpSchema.PaginatedRequest;
 import io.modelcontextprotocol.spec.McpSchema.Root;
-import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler;
-import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
 import io.modelcontextprotocol.util.Assert;
 import io.modelcontextprotocol.util.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -75,6 +75,7 @@ import reactor.core.publisher.Mono;
  * @author Dariusz Jędrzejczyk
  * @author Christian Tzolov
  * @author Jihoon Kim
+ * @author Anurag Pant
  * @see McpClient
  * @see McpSchema
  * @see McpClientSession
@@ -152,16 +153,33 @@ public class McpAsyncClient {
 	 */
 	private final LifecycleInitializer initializer;
 
+	/**
+	 * JSON schema validator to use for validating tool responses against output schemas.
+	 */
+	private final JsonSchemaValidator jsonSchemaValidator;
+
+	/**
+	 * Cached tool output schemas.
+	 */
+	private final ConcurrentHashMap<String, Map<String, Object>> toolsOutputSchemaCache;
+
+	/**
+	 * Whether to enable automatic schema caching during callTool operations.
+	 */
+	private final boolean enableCallToolSchemaCaching;
+
 	/**
 	 * 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 initializationTimeout the max timeout to await for the client-server
-	 * @param features the MCP Client supported features.
+	 * @param jsonSchemaValidator the JSON schema validator to use for validating tool
+	 * @param features the MCP Client supported features. responses against output
+	 * schemas.
 	 */
 	McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout,
-			McpClientFeatures.Async features) {
+			JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features) {
 
 		Assert.notNull(transport, "Transport must not be null");
 		Assert.notNull(requestTimeout, "Request timeout must not be null");
@@ -171,6 +189,9 @@ public class McpAsyncClient {
 		this.clientCapabilities = features.clientCapabilities();
 		this.transport = transport;
 		this.roots = new ConcurrentHashMap<>(features.roots());
+		this.jsonSchemaValidator = jsonSchemaValidator;
+		this.toolsOutputSchemaCache = new ConcurrentHashMap<>();
+		this.enableCallToolSchemaCaching = features.enableCallToolSchemaCaching();
 
 		// Request Handlers
 		Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();
@@ -273,9 +294,30 @@ public class McpAsyncClient {
 		notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_PROGRESS,
 				asyncProgressNotificationHandler(progressConsumersFinal));
 
+		Function<Initialization, Mono<Void>> postInitializationHook = init -> {
+
+			if (init.initializeResult().capabilities().tools() == null || !enableCallToolSchemaCaching) {
+				return Mono.empty();
+			}
+
+			return this.listToolsInternal(init, McpSchema.FIRST_PAGE).doOnNext(listToolsResult -> {
+				listToolsResult.tools()
+					.forEach(tool -> logger.debug("Tool {} schema: {}", tool.name(), tool.outputSchema()));
+				if (enableCallToolSchemaCaching && listToolsResult.tools() != null) {
+					// Cache tools output schema
+					listToolsResult.tools()
+						.stream()
+						.filter(tool -> tool.outputSchema() != null)
+						.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema()));
+				}
+			}).then();
+		};
+
 		this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, transport.protocolVersions(),
 				initializationTimeout, ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
-						notificationHandlers, con -> con.contextWrite(ctx)));
+						notificationHandlers, con -> con.contextWrite(ctx)),
+				postInitializationHook);
+
 		this.transport.setExceptionHandler(this.initializer::handleException);
 	}
 
@@ -387,7 +429,7 @@ public class McpAsyncClient {
 	 * </p>
 	 */
 	public Mono<McpSchema.InitializeResult> initialize() {
-		return this.initializer.withIntitialization("by explicit API call", init -> Mono.just(init.initializeResult()));
+		return this.initializer.withInitialization("by explicit API call", init -> Mono.just(init.initializeResult()));
 	}
 
 	// --------------------------
@@ -399,7 +441,7 @@ public class McpAsyncClient {
 	 * @return A Mono that completes with the server's ping response
 	 */
 	public Mono<Object> ping() {
-		return this.initializer.withIntitialization("pinging the server",
+		return this.initializer.withInitialization("pinging the server",
 				init -> init.mcpSession().sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF));
 	}
 
@@ -480,7 +522,7 @@ public class McpAsyncClient {
 	 * @return A Mono that completes when the notification is sent.
 	 */
 	public Mono<Void> rootsListChangedNotification() {
-		return this.initializer.withIntitialization("sending roots list changed notification",
+		return this.initializer.withInitialization("sending roots list changed notification",
 				init -> init.mcpSession().sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED));
 	}
 
@@ -539,15 +581,45 @@ public class McpAsyncClient {
 	 * @see #listTools()
 	 */
 	public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToolRequest) {
-		return this.initializer.withIntitialization("calling tools", init -> {
+		return this.initializer.withInitialization("calling tool", init -> {
 			if (init.initializeResult().capabilities().tools() == null) {
 				return Mono.error(new IllegalStateException("Server does not provide tools capability"));
 			}
+
 			return init.mcpSession()
-				.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
+				.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF)
+				.flatMap(result -> Mono.just(validateToolResult(callToolRequest.name(), result)));
 		});
 	}
 
+	private McpSchema.CallToolResult validateToolResult(String toolName, McpSchema.CallToolResult result) {
+
+		if (!this.enableCallToolSchemaCaching || result == null || result.isError() == Boolean.TRUE) {
+			// if tool schema caching is disabled or tool call resulted in an error - skip
+			// validation and return the result as it is
+			return result;
+		}
+
+		Map<String, Object> optOutputSchema = this.toolsOutputSchemaCache.get(toolName);
+
+		if (optOutputSchema == null) {
+			logger.warn(
+					"Calling a tool with no outputSchema is not expected to return result with structured content, but got: {}",
+					result.structuredContent());
+			return result;
+		}
+
+		// Validate the tool output against the cached output schema
+		var validation = this.jsonSchemaValidator.validate(optOutputSchema, result.structuredContent());
+
+		if (!validation.valid()) {
+			logger.warn("Tool call result validation failed: {}", validation.errorMessage());
+			throw new IllegalArgumentException("Tool call result validation failed: " + validation.errorMessage());
+		}
+
+		return result;
+	}
+
 	/**
 	 * Retrieves the list of all tools provided by the server.
 	 * @return A Mono that emits the list of all tools result
@@ -568,14 +640,26 @@ public class McpAsyncClient {
 	 * @return A Mono that emits the list of tools result
 	 */
 	public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
-		return this.initializer.withIntitialization("listing tools", init -> {
+		return this.initializer.withInitialization("listing tools", init -> this.listToolsInternal(init, cursor));
-			if (init.initializeResult().capabilities().tools() == null) {
+	}
-				return Mono.error(new IllegalStateException("Server does not provide tools capability"));
+
-			}
+	private Mono<McpSchema.ListToolsResult> listToolsInternal(Initialization init, String cursor) {
-			return init.mcpSession()
+
-				.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
+		if (init.initializeResult().capabilities().tools() == null) {
-						LIST_TOOLS_RESULT_TYPE_REF);
+			return Mono.error(new IllegalStateException("Server does not provide tools capability"));
-		});
+		}
+		return init.mcpSession()
+			.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
+					LIST_TOOLS_RESULT_TYPE_REF)
+			.doOnNext(result -> {
+				if (this.enableCallToolSchemaCaching && result.tools() != null) {
+					// Cache tools output schema
+					result.tools()
+						.stream()
+						.filter(tool -> tool.outputSchema() != null)
+						.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema()));
+				}
+			});
 	}
 
 	private NotificationHandler asyncToolsChangeNotificationHandler(
@@ -632,7 +716,7 @@ public class McpAsyncClient {
 	 * @see #readResource(McpSchema.Resource)
 	 */
 	public Mono<McpSchema.ListResourcesResult> listResources(String cursor) {
-		return this.initializer.withIntitialization("listing resources", init -> {
+		return this.initializer.withInitialization("listing resources", init -> {
 			if (init.initializeResult().capabilities().resources() == null) {
 				return Mono.error(new IllegalStateException("Server does not provide the resources capability"));
 			}
@@ -664,7 +748,7 @@ public class McpAsyncClient {
 	 * @see McpSchema.ReadResourceResult
 	 */
 	public Mono<McpSchema.ReadResourceResult> readResource(McpSchema.ReadResourceRequest readResourceRequest) {
-		return this.initializer.withIntitialization("reading resources", init -> {
+		return this.initializer.withInitialization("reading resources", init -> {
 			if (init.initializeResult().capabilities().resources() == null) {
 				return Mono.error(new IllegalStateException("Server does not provide the resources capability"));
 			}
@@ -702,7 +786,7 @@ public class McpAsyncClient {
 	 * @see McpSchema.ListResourceTemplatesResult
 	 */
 	public Mono<McpSchema.ListResourceTemplatesResult> listResourceTemplates(String cursor) {
-		return this.initializer.withIntitialization("listing resource templates", init -> {
+		return this.initializer.withInitialization("listing resource templates", init -> {
 			if (init.initializeResult().capabilities().resources() == null) {
 				return Mono.error(new IllegalStateException("Server does not provide the resources capability"));
 			}
@@ -722,7 +806,7 @@ public class McpAsyncClient {
 	 * @see #unsubscribeResource(McpSchema.UnsubscribeRequest)
 	 */
 	public Mono<Void> subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
-		return this.initializer.withIntitialization("subscribing to resources", init -> init.mcpSession()
+		return this.initializer.withInitialization("subscribing to resources", init -> init.mcpSession()
 			.sendRequest(McpSchema.METHOD_RESOURCES_SUBSCRIBE, subscribeRequest, VOID_TYPE_REFERENCE));
 	}
 
@@ -736,7 +820,7 @@ public class McpAsyncClient {
 	 * @see #subscribeResource(McpSchema.SubscribeRequest)
 	 */
 	public Mono<Void> unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
-		return this.initializer.withIntitialization("unsubscribing from resources", init -> init.mcpSession()
+		return this.initializer.withInitialization("unsubscribing from resources", init -> init.mcpSession()
 			.sendRequest(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, unsubscribeRequest, VOID_TYPE_REFERENCE));
 	}
 
@@ -802,7 +886,7 @@ public class McpAsyncClient {
 	 * @see #getPrompt(GetPromptRequest)
 	 */
 	public Mono<ListPromptsResult> listPrompts(String cursor) {
-		return this.initializer.withIntitialization("listing prompts", init -> init.mcpSession()
+		return this.initializer.withInitialization("listing prompts", init -> init.mcpSession()
 			.sendRequest(McpSchema.METHOD_PROMPT_LIST, new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF));
 	}
 
@@ -816,7 +900,7 @@ public class McpAsyncClient {
 	 * @see #listPrompts()
 	 */
 	public Mono<GetPromptResult> getPrompt(GetPromptRequest getPromptRequest) {
-		return this.initializer.withIntitialization("getting prompts", init -> init.mcpSession()
+		return this.initializer.withInitialization("getting prompts", init -> init.mcpSession()
 			.sendRequest(McpSchema.METHOD_PROMPT_GET, getPromptRequest, GET_PROMPT_RESULT_TYPE_REF));
 	}
 
@@ -834,14 +918,6 @@ public class McpAsyncClient {
 	// --------------------------
 	// Logging
 	// --------------------------
-	/**
-	 * Create a notification handler for logging notifications from the server. This
-	 * handler automatically distributes logging messages to all registered consumers.
-	 * @param loggingConsumers List of consumers that will be notified when a logging
-	 * message is received. Each consumer receives the logging message notification.
-	 * @return A NotificationHandler that processes log notifications by distributing the
-	 * message to all registered consumers
-	 */
 	private NotificationHandler asyncLoggingNotificationHandler(
 			List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumers) {
 
@@ -867,7 +943,7 @@ public class McpAsyncClient {
 			return Mono.error(new IllegalArgumentException("Logging level must not be null"));
 		}
 
-		return this.initializer.withIntitialization("setting logging level", init -> {
+		return this.initializer.withInitialization("setting logging level", init -> {
 			if (init.initializeResult().capabilities().logging() == null) {
 				return Mono.error(new IllegalStateException("Server's Logging capabilities are not enabled!"));
 			}
@@ -876,15 +952,6 @@ public class McpAsyncClient {
 		});
 	}
 
-	/**
-	 * Create a notification handler for progress notifications from the server. This
-	 * handler automatically distributes progress notifications to all registered
-	 * consumers.
-	 * @param progressConsumers List of consumers that will be notified when a progress
-	 * message is received. Each consumer receives the progress notification.
-	 * @return A NotificationHandler that processes progress notifications by distributing
-	 * the message to all registered consumers
-	 */
 	private NotificationHandler asyncProgressNotificationHandler(
 			List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers) {
 
@@ -924,7 +991,7 @@ public class McpAsyncClient {
 	 * @see McpSchema.CompleteResult
 	 */
 	public Mono<McpSchema.CompleteResult> completeCompletion(McpSchema.CompleteRequest completeRequest) {
-		return this.initializer.withIntitialization("complete completions", init -> init.mcpSession()
+		return this.initializer.withInitialization("complete completions", init -> init.mcpSession()
 			.sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF));
 	}
 

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

@@ -13,6 +13,7 @@ import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
+import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
@@ -99,6 +100,7 @@ import reactor.core.publisher.Mono;
  *
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
+ * @author Anurag Pant
  * @see McpAsyncClient
  * @see McpSyncClient
  * @see McpTransport
@@ -187,6 +189,10 @@ public interface McpClient {
 
 		private Supplier<McpTransportContext> contextProvider = () -> McpTransportContext.EMPTY;
 
+		private JsonSchemaValidator jsonSchemaValidator;
+
+		private boolean enableCallToolSchemaCaching = false; // Default to false
+
 		private SyncSpec(McpClientTransport transport) {
 			Assert.notNull(transport, "Transport must not be null");
 			this.transport = transport;
@@ -429,6 +435,32 @@ public interface McpClient {
 			return this;
 		}
 
+		/**
+		 * Add a {@link JsonSchemaValidator} to validate the JSON structure of the
+		 * structured output.
+		 * @param jsonSchemaValidator A validator to validate the JSON structure of the
+		 * structured output. Must not be null.
+		 * @return This builder for method chaining
+		 * @throws IllegalArgumentException if jsonSchemaValidator is null
+		 */
+		public SyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
+			Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
+			this.jsonSchemaValidator = jsonSchemaValidator;
+			return this;
+		}
+
+		/**
+		 * Enables automatic schema caching during callTool operations. When a tool's
+		 * output schema is not found in the cache, callTool will automatically fetch and
+		 * cache all tool schemas via listTools.
+		 * @param enableCallToolSchemaCaching true to enable, false to disable
+		 * @return This builder instance for method chaining
+		 */
+		public SyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) {
+			this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
+			return this;
+		}
+
 		/**
 		 * Create an instance of {@link McpSyncClient} with the provided configurations or
 		 * sensible defaults.
@@ -438,13 +470,13 @@ public interface McpClient {
 			McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities,
 					this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
 					this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, this.samplingHandler,
-					this.elicitationHandler);
+					this.elicitationHandler, this.enableCallToolSchemaCaching);
 
 			McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
 
-			return new McpSyncClient(
+			return new McpSyncClient(new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout,
-					new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures),
+					jsonSchemaValidator != null ? jsonSchemaValidator : JsonSchemaValidator.getDefault(),
-					this.contextProvider);
+					asyncFeatures), this.contextProvider);
 		}
 
 	}
@@ -495,6 +527,10 @@ public interface McpClient {
 
 		private Function<ElicitRequest, Mono<ElicitResult>> elicitationHandler;
 
+		private JsonSchemaValidator jsonSchemaValidator;
+
+		private boolean enableCallToolSchemaCaching = false; // Default to false
+
 		private AsyncSpec(McpClientTransport transport) {
 			Assert.notNull(transport, "Transport must not be null");
 			this.transport = transport;
@@ -741,17 +777,45 @@ public interface McpClient {
 			return this;
 		}
 
+		/**
+		 * Sets the JSON schema validator to use for validating tool responses against
+		 * output schemas.
+		 * @param jsonSchemaValidator The validator to use. Must not be null.
+		 * @return This builder instance for method chaining
+		 * @throws IllegalArgumentException if jsonSchemaValidator is null
+		 */
+		public AsyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
+			Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
+			this.jsonSchemaValidator = jsonSchemaValidator;
+			return this;
+		}
+
+		/**
+		 * Enables automatic schema caching during callTool operations. When a tool's
+		 * output schema is not found in the cache, callTool will automatically fetch and
+		 * cache all tool schemas via listTools.
+		 * @param enableCallToolSchemaCaching true to enable, false to disable
+		 * @return This builder instance for method chaining
+		 */
+		public AsyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) {
+			this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
+			return this;
+		}
+
 		/**
 		 * Create an instance of {@link McpAsyncClient} with the provided configurations
 		 * or sensible defaults.
 		 * @return a new instance of {@link McpAsyncClient}.
 		 */
 		public McpAsyncClient build() {
+			var jsonSchemaValidator = (this.jsonSchemaValidator != null) ? this.jsonSchemaValidator
+					: JsonSchemaValidator.getDefault();
 			return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout,
+					jsonSchemaValidator,
 					new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
 							this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
 							this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
-							this.samplingHandler, this.elicitationHandler));
+							this.samplingHandler, this.elicitationHandler, this.enableCallToolSchemaCaching));
 		}
 
 	}

mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java

@@ -62,6 +62,7 @@ class McpClientFeatures {
 	 * @param progressConsumers the progress consumers.
 	 * @param samplingHandler the sampling handler.
 	 * @param elicitationHandler the elicitation handler.
+	 * @param enableCallToolSchemaCaching whether to enable call tool schema caching.
 	 */
 	record Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
 			Map<String, McpSchema.Root> roots, List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers,
@@ -71,7 +72,8 @@ class McpClientFeatures {
 			List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
 			List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers,
 			Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
-			Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
+			Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler,
+			boolean enableCallToolSchemaCaching) {
 
 		/**
 		 * Create an instance and validate the arguments.
@@ -84,6 +86,7 @@ class McpClientFeatures {
 		 * @param progressConsumers the progress consumers.
 		 * @param samplingHandler the sampling handler.
 		 * @param elicitationHandler the elicitation handler.
+		 * @param enableCallToolSchemaCaching whether to enable call tool schema caching.
 		 */
 		public Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
 				Map<String, McpSchema.Root> roots,
@@ -94,7 +97,8 @@ class McpClientFeatures {
 				List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
 				List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers,
 				Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
-				Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
+				Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler,
+				boolean enableCallToolSchemaCaching) {
 
 			Assert.notNull(clientInfo, "Client info must not be null");
 			this.clientInfo = clientInfo;
@@ -113,6 +117,7 @@ class McpClientFeatures {
 			this.progressConsumers = progressConsumers != null ? progressConsumers : List.of();
 			this.samplingHandler = samplingHandler;
 			this.elicitationHandler = elicitationHandler;
+			this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
 		}
 
 		/**
@@ -129,7 +134,7 @@ class McpClientFeatures {
 				Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
 			this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers,
 					resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler,
-					elicitationHandler);
+					elicitationHandler, false);
 		}
 
 		/**
@@ -187,7 +192,8 @@ class McpClientFeatures {
 
 			return new Async(syncSpec.clientInfo(), syncSpec.clientCapabilities(), syncSpec.roots(),
 					toolsChangeConsumers, resourcesChangeConsumers, resourcesUpdateConsumers, promptsChangeConsumers,
-					loggingConsumers, progressConsumers, samplingHandler, elicitationHandler);
+					loggingConsumers, progressConsumers, samplingHandler, elicitationHandler,
+					syncSpec.enableCallToolSchemaCaching);
 		}
 	}
 
@@ -205,6 +211,7 @@ class McpClientFeatures {
 	 * @param progressConsumers the progress consumers.
 	 * @param samplingHandler the sampling handler.
 	 * @param elicitationHandler the elicitation handler.
+	 * @param enableCallToolSchemaCaching whether to enable call tool schema caching.
 	 */
 	public record Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
 			Map<String, McpSchema.Root> roots, List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers,
@@ -214,7 +221,8 @@ class McpClientFeatures {
 			List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
 			List<Consumer<McpSchema.ProgressNotification>> progressConsumers,
 			Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
-			Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
+			Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler,
+			boolean enableCallToolSchemaCaching) {
 
 		/**
 		 * Create an instance and validate the arguments.
@@ -229,6 +237,7 @@ class McpClientFeatures {
 		 * @param progressConsumers the progress consumers.
 		 * @param samplingHandler the sampling handler.
 		 * @param elicitationHandler the elicitation handler.
+		 * @param enableCallToolSchemaCaching whether to enable call tool schema caching.
 		 */
 		public Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
 				Map<String, McpSchema.Root> roots, List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers,
@@ -238,7 +247,8 @@ class McpClientFeatures {
 				List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
 				List<Consumer<McpSchema.ProgressNotification>> progressConsumers,
 				Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
-				Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
+				Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler,
+				boolean enableCallToolSchemaCaching) {
 
 			Assert.notNull(clientInfo, "Client info must not be null");
 			this.clientInfo = clientInfo;
@@ -257,6 +267,7 @@ class McpClientFeatures {
 			this.progressConsumers = progressConsumers != null ? progressConsumers : List.of();
 			this.samplingHandler = samplingHandler;
 			this.elicitationHandler = elicitationHandler;
+			this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
 		}
 
 		/**
@@ -272,7 +283,7 @@ class McpClientFeatures {
 				Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
 			this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers,
 					resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler,
-					elicitationHandler);
+					elicitationHandler, false);
 		}
 	}
 

mcp-core/src/main/java/io/modelcontextprotocol/spec/McpSchema.java

@@ -17,7 +17,6 @@ import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
-import com.fasterxml.jackson.annotation.JsonTypeInfo.As;
 import io.modelcontextprotocol.json.McpJsonMapper;
 import io.modelcontextprotocol.json.TypeRef;
 import io.modelcontextprotocol.util.Assert;

mcp-core/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java

@@ -4,15 +4,14 @@
 
 package io.modelcontextprotocol.client;
 
+import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
+import io.modelcontextprotocol.spec.McpClientTransport;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Timeout;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.strategy.Wait;
 
-import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
-import io.modelcontextprotocol.spec.McpClientTransport;
-
 @Timeout(15)
 public class HttpClientStreamableHttpAsyncClientTests extends AbstractMcpAsyncClientTests {
 

mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerPostInitializationHookTests.java

@@ -0,0 +1,280 @@
+/*
+ * Copyright 2024-2024 the original author or authors.
+ */
+
+package io.modelcontextprotocol.client;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+
+import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
+import io.modelcontextprotocol.spec.McpClientSession;
+import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+import reactor.test.StepVerifier;
+import reactor.util.context.ContextView;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link LifecycleInitializer} postInitializationHook functionality.
+ * 
+ * @author Christian Tzolov
+ */
+class LifecycleInitializerPostInitializationHookTests {
+
+	private static final Duration INITIALIZATION_TIMEOUT = Duration.ofSeconds(5);
+
+	private static final McpSchema.ClientCapabilities CLIENT_CAPABILITIES = McpSchema.ClientCapabilities.builder()
+		.build();
+
+	private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0");
+
+	private static final List<String> PROTOCOL_VERSIONS = List.of("1.0.0", "2.0.0");
+
+	private static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult("2.0.0",
+			McpSchema.ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"),
+			"Test instructions");
+
+	@Mock
+	private McpClientSession mockClientSession;
+
+	@Mock
+	private Function<ContextView, McpClientSession> mockSessionSupplier;
+
+	@Mock
+	private Function<Initialization, Mono<Void>> mockPostInitializationHook;
+
+	private LifecycleInitializer initializer;
+
+	@BeforeEach
+	void setUp() {
+		MockitoAnnotations.openMocks(this);
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty());
+		when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession);
+		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
+			.thenReturn(Mono.just(MOCK_INIT_RESULT));
+		when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
+			.thenReturn(Mono.empty());
+		when(mockClientSession.closeGracefully()).thenReturn(Mono.empty());
+
+		initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
+				INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
+	}
+
+	@Test
+	void shouldInvokePostInitializationHook() {
+		AtomicReference<Initialization> capturedInit = new AtomicReference<>();
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
+			capturedInit.set(invocation.getArgument(0));
+			return Mono.empty();
+		});
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		// Verify hook was called
+		verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
+
+		// Verify the hook received correct initialization data
+		assertThat(capturedInit.get()).isNotNull();
+		assertThat(capturedInit.get().mcpSession()).isEqualTo(mockClientSession);
+		assertThat(capturedInit.get().initializeResult()).isEqualTo(MOCK_INIT_RESULT);
+	}
+
+	@Test
+	void shouldInvokePostInitializationHookOnlyOnce() {
+		// First initialization
+		StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
+			.expectNext("result1")
+			.verifyComplete();
+
+		// Second call should reuse initialization and NOT call hook again
+		StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
+			.expectNext("result2")
+			.verifyComplete();
+
+		// Hook should only be called once
+		verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
+	}
+
+	@Test
+	void shouldInvokePostInitializationHookOnlyOnceWithConcurrentRequests() {
+		AtomicInteger hookInvocationCount = new AtomicInteger(0);
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
+			hookInvocationCount.incrementAndGet();
+			return Mono.empty();
+		});
+
+		// Start multiple concurrent initializations
+		Mono<String> init1 = initializer.withInitialization("test1", init -> Mono.just("result1"))
+			.subscribeOn(Schedulers.parallel());
+		Mono<String> init2 = initializer.withInitialization("test2", init -> Mono.just("result2"))
+			.subscribeOn(Schedulers.parallel());
+		Mono<String> init3 = initializer.withInitialization("test3", init -> Mono.just("result3"))
+			.subscribeOn(Schedulers.parallel());
+
+		// TODO: can we assume the order of results?
+		StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> {
+			assertThat(tuple.getT1()).isEqualTo("result1");
+			assertThat(tuple.getT2()).isEqualTo("result2");
+			assertThat(tuple.getT3()).isEqualTo("result3");
+		}).verifyComplete();
+
+		// Hook should only be called once despite concurrent requests
+		assertThat(hookInvocationCount.get()).isEqualTo(1);
+	}
+
+	@Test
+	void shouldFailInitializationWhenPostInitializationHookFails() {
+		RuntimeException hookError = new RuntimeException("Post-initialization hook failed");
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.error(hookError));
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectErrorMatches(ex -> ex instanceof RuntimeException && ex.getCause() == hookError)
+			.verify();
+
+		// Verify initialization was not completed
+		assertThat(initializer.isInitialized()).isFalse();
+		assertThat(initializer.currentInitializationResult()).isNull();
+
+		// Verify the hook was called
+		verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
+	}
+
+	@Test
+	void shouldNotInvokePostInitializationHookWhenInitializationFails() {
+		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
+			.thenReturn(Mono.error(new RuntimeException("Initialization failed")));
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectError(RuntimeException.class)
+			.verify();
+
+		// Hook should NOT be called when initialization fails
+		verify(mockPostInitializationHook, never()).apply(any(Initialization.class));
+	}
+
+	@Test
+	void shouldNotInvokePostInitializationHookWhenNotificationFails() {
+		when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
+			.thenReturn(Mono.error(new RuntimeException("Notification failed")));
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectError(RuntimeException.class)
+			.verify();
+
+		// Hook should NOT be called when notification fails
+		verify(mockPostInitializationHook, never()).apply(any(Initialization.class));
+	}
+
+	@Test
+	void shouldInvokePostInitializationHookAgainAfterReinitialization() {
+		AtomicInteger hookInvocationCount = new AtomicInteger(0);
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
+			hookInvocationCount.incrementAndGet();
+			return Mono.empty();
+		});
+
+		// First initialization
+		StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
+			.expectNext("result1")
+			.verifyComplete();
+
+		assertThat(hookInvocationCount.get()).isEqualTo(1);
+
+		// Simulate transport session exception to trigger re-initialization
+		initializer.handleException(new McpTransportSessionNotFoundException("Session lost"));
+
+		// Hook should be called twice (once for each initialization)
+		assertThat(hookInvocationCount.get()).isEqualTo(2);
+	}
+
+	@Test
+	void shouldAllowPostInitializationHookToPerformAsyncOperations() {
+		AtomicInteger operationCount = new AtomicInteger(0);
+
+		when(mockPostInitializationHook.apply(any(Initialization.class)))
+			.thenReturn(Mono.fromRunnable(() -> operationCount.incrementAndGet()).then());
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		// Verify the async operation was executed
+		assertThat(operationCount.get()).isEqualTo(1);
+		verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
+	}
+
+	@Test
+	void shouldProvideCorrectInitializationDataToHook() {
+		AtomicReference<McpClientSession> capturedSession = new AtomicReference<>();
+		AtomicReference<McpSchema.InitializeResult> capturedResult = new AtomicReference<>();
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
+			Initialization init = invocation.getArgument(0);
+			capturedSession.set(init.mcpSession());
+			capturedResult.set(init.initializeResult());
+			return Mono.empty();
+		});
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		// Verify the hook received the correct session and result
+		assertThat(capturedSession.get()).isEqualTo(mockClientSession);
+		assertThat(capturedResult.get()).isEqualTo(MOCK_INIT_RESULT);
+		assertThat(capturedResult.get().protocolVersion()).isEqualTo("2.0.0");
+		assertThat(capturedResult.get().serverInfo().name()).isEqualTo("test-server");
+	}
+
+	@Test
+	void shouldInvokePostInitializationHookAfterSuccessfulInitialization() {
+		AtomicReference<Boolean> notificationSent = new AtomicReference<>(false);
+		AtomicReference<Boolean> hookCalledAfterNotification = new AtomicReference<>(false);
+
+		when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
+			.thenAnswer(invocation -> {
+				notificationSent.set(true);
+				return Mono.empty();
+			});
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
+			// Due to flatMap chaining in doInitialize, if the hook is called,
+			// the notification must have been sent first
+			hookCalledAfterNotification.set(notificationSent.get());
+			return Mono.empty();
+		});
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		// Verify the hook was called and notification was already sent at that point
+		assertThat(hookCalledAfterNotification.get()).isTrue();
+		verify(mockClientSession).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any());
+		verify(mockPostInitializationHook).apply(any(Initialization.class));
+	}
+
+}

mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java

@@ -10,14 +10,14 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
+import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
+import io.modelcontextprotocol.spec.McpClientSession;
+import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-
-import io.modelcontextprotocol.spec.McpClientSession;
-import io.modelcontextprotocol.spec.McpSchema;
-import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
@@ -58,12 +58,16 @@ class LifecycleInitializerTests {
 	@Mock
 	private Function<ContextView, McpClientSession> mockSessionSupplier;
 
+	@Mock
+	private Function<Initialization, Mono<Void>> mockPostInitializationHook;
+
 	private LifecycleInitializer initializer;
 
 	@BeforeEach
 	void setUp() {
 		MockitoAnnotations.openMocks(this);
 
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty());
 		when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession);
 		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
 			.thenReturn(Mono.just(MOCK_INIT_RESULT));
@@ -72,45 +76,45 @@ class LifecycleInitializerTests {
 		when(mockClientSession.closeGracefully()).thenReturn(Mono.empty());
 
 		initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
-				INITIALIZATION_TIMEOUT, mockSessionSupplier);
+				INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
 	}
 
 	@Test
 	void constructorShouldValidateParameters() {
 		assertThatThrownBy(() -> new LifecycleInitializer(null, CLIENT_INFO, PROTOCOL_VERSIONS, INITIALIZATION_TIMEOUT,
-				mockSessionSupplier))
+				mockSessionSupplier, mockPostInitializationHook))
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("Client capabilities must not be null");
 
 		assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, null, PROTOCOL_VERSIONS,
-				INITIALIZATION_TIMEOUT, mockSessionSupplier))
+				INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("Client info must not be null");
 
 		assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, null,
-				INITIALIZATION_TIMEOUT, mockSessionSupplier))
+				INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("Protocol versions must not be empty");
 
 		assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, List.of(),
-				INITIALIZATION_TIMEOUT, mockSessionSupplier))
+				INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("Protocol versions must not be empty");
 
 		assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, null,
-				mockSessionSupplier))
+				mockSessionSupplier, mockPostInitializationHook))
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("Initialization timeout must not be null");
 
 		assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
-				INITIALIZATION_TIMEOUT, null))
+				INITIALIZATION_TIMEOUT, null, mockPostInitializationHook))
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("Session supplier must not be null");
 	}
 
 	@Test
 	void shouldInitializeSuccessfully() {
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.assertNext(result -> {
 				assertThat(result).isEqualTo(MOCK_INIT_RESULT);
 				assertThat(initializer.isInitialized()).isTrue();
@@ -132,7 +136,7 @@ class LifecycleInitializerTests {
 			return Mono.just(MOCK_INIT_RESULT);
 		});
 
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.assertNext(result -> {
 				assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0"); // Latest
 																						// version
@@ -152,7 +156,7 @@ class LifecycleInitializerTests {
 		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
 			.thenReturn(Mono.just(unsupportedResult));
 
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectError(RuntimeException.class)
 			.verify();
 
@@ -167,13 +171,13 @@ class LifecycleInitializerTests {
 		Duration SLOW_RESPONSE_DELAY = Duration.ofSeconds(5);
 
 		LifecycleInitializer shortTimeoutInitializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO,
-				PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier);
+				PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
 
 		when(mockClientSession.<McpSchema.InitializeResult>sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
 			.thenReturn(Mono.just(MOCK_INIT_RESULT).delayElement(SLOW_RESPONSE_DELAY, virtualTimeScheduler));
 
 		StepVerifier
-			.withVirtualTime(() -> shortTimeoutInitializer.withIntitialization("test",
+			.withVirtualTime(() -> shortTimeoutInitializer.withInitialization("test",
 					init -> Mono.just(init.initializeResult())), () -> virtualTimeScheduler, Long.MAX_VALUE)
 			.expectSubscription()
 			.expectNoEvent(INITIALIZE_TIMEOUT)
@@ -184,12 +188,12 @@ class LifecycleInitializerTests {
 	@Test
 	void shouldReuseExistingInitialization() {
 		// First initialization
-		StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1")))
+		StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
 			.expectNext("result1")
 			.verifyComplete();
 
 		// Second call should reuse the same initialization
-		StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2")))
+		StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
 			.expectNext("result2")
 			.verifyComplete();
 
@@ -209,11 +213,11 @@ class LifecycleInitializerTests {
 
 		// Start multiple concurrent initializations using subscribeOn with parallel
 		// scheduler
-		Mono<String> init1 = initializer.withIntitialization("test1", init -> Mono.just("result1"))
+		Mono<String> init1 = initializer.withInitialization("test1", init -> Mono.just("result1"))
 			.subscribeOn(Schedulers.parallel());
-		Mono<String> init2 = initializer.withIntitialization("test2", init -> Mono.just("result2"))
+		Mono<String> init2 = initializer.withInitialization("test2", init -> Mono.just("result2"))
 			.subscribeOn(Schedulers.parallel());
-		Mono<String> init3 = initializer.withIntitialization("test3", init -> Mono.just("result3"))
+		Mono<String> init3 = initializer.withInitialization("test3", init -> Mono.just("result3"))
 			.subscribeOn(Schedulers.parallel());
 
 		StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> {
@@ -235,7 +239,7 @@ class LifecycleInitializerTests {
 			// succeeds on the second call
 			.thenReturn(Mono.just(MOCK_INIT_RESULT));
 
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectError(RuntimeException.class)
 			.verify();
 
@@ -244,7 +248,7 @@ class LifecycleInitializerTests {
 
 		// The initializer can recover from previous errors
 		StepVerifier
-			.create(initializer.withIntitialization("successful init", init -> Mono.just(init.initializeResult())))
+			.create(initializer.withInitialization("successful init", init -> Mono.just(init.initializeResult())))
 			.expectNext(MOCK_INIT_RESULT)
 			.verifyComplete();
 
@@ -255,7 +259,7 @@ class LifecycleInitializerTests {
 	@Test
 	void shouldHandleTransportSessionNotFoundException() {
 		// successful initialization first
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectNext(MOCK_INIT_RESULT)
 			.verifyComplete();
 
@@ -277,7 +281,7 @@ class LifecycleInitializerTests {
 	@Test
 	void shouldHandleOtherExceptions() {
 		// Simulate a successful initialization first
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectNext(MOCK_INIT_RESULT)
 			.verifyComplete();
 
@@ -295,7 +299,7 @@ class LifecycleInitializerTests {
 
 	@Test
 	void shouldCloseGracefully() {
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectNext(MOCK_INIT_RESULT)
 			.verifyComplete();
 
@@ -307,7 +311,7 @@ class LifecycleInitializerTests {
 
 	@Test
 	void shouldCloseImmediately() {
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectNext(MOCK_INIT_RESULT)
 			.verifyComplete();
 
@@ -342,7 +346,7 @@ class LifecycleInitializerTests {
 					new McpSchema.Implementation("test-server", "1.0.0"), "Test instructions"));
 		});
 
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.assertNext(result -> {
 				// Latest from new versions
 				assertThat(capturedRequest.get().protocolVersion()).isEqualTo("4.0.0");
@@ -363,7 +367,7 @@ class LifecycleInitializerTests {
 		});
 
 		StepVerifier
-			.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))
+			.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))
 				.contextWrite(Context.of(contextKey, contextValue)))
 			.expectNext(MOCK_INIT_RESULT)
 			.verifyComplete();
@@ -374,7 +378,7 @@ class LifecycleInitializerTests {
 
 	@Test
 	void shouldProvideAccessToMcpSessionAndInitializeResult() {
-		StepVerifier.create(initializer.withIntitialization("test", init -> {
+		StepVerifier.create(initializer.withInitialization("test", init -> {
 			assertThat(init.mcpSession()).isEqualTo(mockClientSession);
 			assertThat(init.initializeResult()).isEqualTo(MOCK_INIT_RESULT);
 			return Mono.just("success");
@@ -386,7 +390,7 @@ class LifecycleInitializerTests {
 		when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
 			.thenReturn(Mono.error(new RuntimeException("Notification failed")));
 
-		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectError(RuntimeException.class)
 			.verify();
 
@@ -403,7 +407,7 @@ class LifecycleInitializerTests {
 	@Test
 	void shouldReinitializeAfterTransportSessionException() {
 		// First initialization
-		StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1")))
+		StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
 			.expectNext("result1")
 			.verifyComplete();
 
@@ -411,7 +415,7 @@ class LifecycleInitializerTests {
 		initializer.handleException(new McpTransportSessionNotFoundException("Session lost"));
 
 		// Should be able to initialize again
-		StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2")))
+		StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
 			.expectNext("result2")
 			.verifyComplete();
 

mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java

@@ -10,12 +10,20 @@ import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.ProtocolVersions;
 
 import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Mono;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
 import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 
 class McpAsyncClientTests {
@@ -24,6 +32,7 @@ class McpAsyncClientTests {
 			"1.0.0");
 
 	public static final McpSchema.ServerCapabilities MOCK_SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder()
+		.tools(true)
 		.build();
 
 	public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult(
@@ -31,6 +40,91 @@ class McpAsyncClientTests {
 
 	private static final String CONTEXT_KEY = "context.key";
 
+	private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput)
+			throws JsonProcessingException {
+
+		// Create tool with or without output schema
+		Map<String, Object> inputSchemaMap = Map.of("type", "object", "properties",
+				Map.of("expression", Map.of("type", "string")), "required", List.of("expression"));
+
+		McpSchema.JsonSchema inputSchema = new McpSchema.JsonSchema("object", inputSchemaMap, null, null, null, null);
+		McpSchema.Tool.Builder toolBuilder = McpSchema.Tool.builder()
+			.name("calculator")
+			.description("Performs mathematical calculations")
+			.inputSchema(inputSchema);
+
+		if (hasOutputSchema) {
+			Map<String, Object> outputSchema = Map.of("type", "object", "properties",
+					Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required",
+					List.of("result", "operation"));
+			toolBuilder.outputSchema(outputSchema);
+		}
+
+		McpSchema.Tool calculatorTool = toolBuilder.build();
+		McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(calculatorTool), null);
+
+		// Create call tool result - valid or invalid based on parameter
+		Map<String, Object> structuredContent = invalidOutput ? Map.of("result", "5", "operation", "add")
+				: Map.of("result", 5, "operation", "add");
+
+		McpSchema.CallToolResult mockCallToolResult = McpSchema.CallToolResult.builder()
+			.addTextContent("Calculation result")
+			.structuredContent(structuredContent)
+			.build();
+
+		return new McpClientTransport() {
+			Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler;
+
+			@Override
+			public Mono<Void> connect(
+					Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
+				this.handler = handler;
+				return Mono.empty();
+			}
+
+			@Override
+			public Mono<Void> closeGracefully() {
+				return Mono.empty();
+			}
+
+			@Override
+			public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
+				if (!(message instanceof McpSchema.JSONRPCRequest request)) {
+					return Mono.empty();
+				}
+
+				McpSchema.JSONRPCResponse response;
+				if (McpSchema.METHOD_INITIALIZE.equals(request.method())) {
+					response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), MOCK_INIT_RESULT,
+							null);
+				}
+				else if (McpSchema.METHOD_TOOLS_LIST.equals(request.method())) {
+					response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), mockToolsResult,
+							null);
+				}
+				else if (McpSchema.METHOD_TOOLS_CALL.equals(request.method())) {
+					response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
+							mockCallToolResult, null);
+				}
+				else {
+					return Mono.empty();
+				}
+
+				return handler.apply(Mono.just(response)).then();
+			}
+
+			@Override
+			public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
+				return JSON_MAPPER.convertValue(data, new TypeRef<>() {
+					@Override
+					public java.lang.reflect.Type getType() {
+						return typeRef.getType();
+					}
+				});
+			}
+		};
+	}
+
 	@Test
 	void validateContextPassedToTransportConnect() {
 		McpClientTransport transport = new McpClientTransport() {
@@ -87,4 +181,64 @@ class McpAsyncClientTests {
 		}).doesNotThrowAnyException();
 	}
 
+	@Test
+	void testCallToolWithOutputSchemaValidationSuccess() throws JsonProcessingException {
+		McpClientTransport transport = createMockTransportForToolValidation(true, false);
+
+		McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
+
+		StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+		StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
+			.expectNextMatches(response -> {
+				assertThat(response).isNotNull();
+				assertThat(response.isError()).isFalse();
+				assertThat(response.structuredContent()).isInstanceOf(Map.class);
+				assertThat((Map<?, ?>) response.structuredContent()).hasSize(2);
+				assertThat(response.content()).hasSize(1);
+				return true;
+			})
+			.verifyComplete();
+
+		StepVerifier.create(client.closeGracefully()).verifyComplete();
+	}
+
+	@Test
+	void testCallToolWithNoOutputSchemaSuccess() throws JsonProcessingException {
+		McpClientTransport transport = createMockTransportForToolValidation(false, false);
+
+		McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
+
+		StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+		StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
+			.expectNextMatches(response -> {
+				assertThat(response).isNotNull();
+				assertThat(response.isError()).isFalse();
+				assertThat(response.structuredContent()).isInstanceOf(Map.class);
+				assertThat((Map<?, ?>) response.structuredContent()).hasSize(2);
+				assertThat(response.content()).hasSize(1);
+				return true;
+			})
+			.verifyComplete();
+
+		StepVerifier.create(client.closeGracefully()).verifyComplete();
+	}
+
+	@Test
+	void testCallToolWithOutputSchemaValidationFailure() throws JsonProcessingException {
+		McpClientTransport transport = createMockTransportForToolValidation(true, true);
+
+		McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
+
+		StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+		StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
+			.expectErrorMatches(ex -> ex instanceof IllegalArgumentException
+					&& ex.getMessage().contains("Tool call result validation failed"))
+			.verify();
+
+		StepVerifier.create(client.closeGracefully()).verifyComplete();
+	}
+
 }

mcp-core/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java

@@ -8,9 +8,9 @@ import java.time.Duration;
 import java.util.List;
 
 import io.modelcontextprotocol.MockMcpClientTransport;
-import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
+import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -22,7 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class McpClientProtocolVersionTests {
 
-	private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(30);
+	private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(300);
 
 	private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0");
 
@@ -46,13 +46,12 @@ class McpClientProtocolVersionTests {
 				assertThat(initRequest.protocolVersion()).isEqualTo(transport.protocolVersions().get(0));
 
 				transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
-						new McpSchema.InitializeResult(protocolVersion, null,
+						new McpSchema.InitializeResult(protocolVersion, ServerCapabilities.builder().build(),
 								new McpSchema.Implementation("test-server", "1.0.0"), null),
 						null));
 			}).assertNext(result -> {
 				assertThat(result.protocolVersion()).isEqualTo(protocolVersion);
 			}).verifyComplete();
-
 		}
 		finally {
 			// Ensure cleanup happens even if test fails
@@ -81,7 +80,7 @@ class McpClientProtocolVersionTests {
 				assertThat(initRequest.protocolVersion()).isIn(List.of(oldVersion, McpSchema.LATEST_PROTOCOL_VERSION));
 
 				transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
-						new McpSchema.InitializeResult(oldVersion, null,
+						new McpSchema.InitializeResult(oldVersion, ServerCapabilities.builder().build(),
 								new McpSchema.Implementation("test-server", "1.0.0"), null),
 						null));
 			}).assertNext(result -> {
@@ -110,7 +109,7 @@ class McpClientProtocolVersionTests {
 				assertThat(request.params()).isInstanceOf(McpSchema.InitializeRequest.class);
 
 				transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
-						new McpSchema.InitializeResult(unsupportedVersion, null,
+						new McpSchema.InitializeResult(unsupportedVersion, ServerCapabilities.builder().build(),
 								new McpSchema.Implementation("test-server", "1.0.0"), null),
 						null));
 			}).expectError(RuntimeException.class).verify();
@@ -143,7 +142,7 @@ class McpClientProtocolVersionTests {
 				assertThat(initRequest.protocolVersion()).isEqualTo(latestVersion);
 
 				transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
-						new McpSchema.InitializeResult(latestVersion, null,
+						new McpSchema.InitializeResult(latestVersion, ServerCapabilities.builder().build(),
 								new McpSchema.Implementation("test-server", "1.0.0"), null),
 						null));
 			}).assertNext(result -> {