Store optional json schemas in cache

Author: pantanurag555

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

@@ -8,6 +8,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -26,6 +27,7 @@
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
 import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
 import io.modelcontextprotocol.spec.McpSchema.GetPromptResult;
+import io.modelcontextprotocol.spec.McpSchema.JsonSchema;
 import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
 import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
 import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
@@ -130,7 +132,7 @@ public class McpAsyncClient {
 	/**
 	 * Cached tool output schemas.
 	 */
-	private final ConcurrentHashMap<String, McpSchema.JsonSchema> toolsOutputSchemaCache;
+	private final ConcurrentHashMap<String, Optional<JsonSchema>> toolsOutputSchemaCache;
 
 	/**
 	 * Roots define the boundaries of where servers can operate within the filesystem,
@@ -301,7 +303,7 @@ public McpSchema.Implementation getClientInfo() {
 	 * Get the cached tool output schemas.
 	 * @return The cached tool output schemas
 	 */
-	public HashMap<String, McpSchema.JsonSchema> getToolsOutputSchemaCache() {
+	public ConcurrentHashMap<String, Optional<JsonSchema>> getToolsOutputSchemaCache() {
 		return this.toolsOutputSchemaCache;
 	}
 
@@ -579,7 +581,8 @@ public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
 					if (result.tools() != null) {
 						// Cache tools output schema
 						result.tools()
-							.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema()));
+							.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(),
+									Optional.ofNullable(tool.outputSchema())));
 					}
 				});
 		});

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

@@ -5,8 +5,9 @@
 package io.modelcontextprotocol.client;
 
 import java.time.Duration;
-import java.util.HashMap;
+import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -244,18 +245,19 @@ public Object ping() {
 	 */
 	public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) {
 		McpSchema.CallToolResult result = this.delegate.callTool(callToolRequest).block();
-		HashMap<String, McpSchema.JsonSchema> toolsOutputSchemaCache = this.delegate.getToolsOutputSchemaCache();
+		ConcurrentHashMap<String, Optional<McpSchema.JsonSchema>> toolsOutputSchemaCache = this.delegate
+			.getToolsOutputSchemaCache();
 		// Should not be triggered but added for completeness
 		if (!toolsOutputSchemaCache.containsKey(callToolRequest.name())) {
 			throw new McpError("Tool with name '" + callToolRequest.name() + "' not found");
 		}
-		if (result != null && toolsOutputSchemaCache.get(callToolRequest.name()) != null) {
+		Optional<McpSchema.JsonSchema> optOutputSchema = toolsOutputSchemaCache.get(callToolRequest.name());
+		if (result != null && optOutputSchema != null && optOutputSchema.isPresent()) {
 			if (result.structuredContent() == null) {
 				throw new McpError("CallToolResult validation failed: structuredContent is null and "
 						+ "does not match tool outputSchema.");
 			}
-
-			McpSchema.JsonSchema outputSchema = toolsOutputSchemaCache.get(callToolRequest.name());
+			McpSchema.JsonSchema outputSchema = optOutputSchema.get();
 
 			try {
 				// Convert outputSchema to string
@@ -284,9 +286,9 @@ public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolReque
 				}
 			}
 			catch (JsonProcessingException e) {
-				// Log error if output schema can't be parsed to prevent erroring out for
-				// successful call tool request
-				logger.error("Encountered exception when parsing outputSchema: {}", e);
+				// Log warning if output schema can't be parsed to prevent erroring out
+				// for successful call tool request
+				logger.warn("Failed to validate CallToolResult: Error parsing tool outputSchema: {}", e);
 			}
 		}
 		return result;