DefaultJsonSchemaValidator.java

/*
 * Copyright 2024-2024 the original author or authors.
 */

package io.modelcontextprotocol.spec;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.networknt.schema.JsonSchema;
import com.networknt.schema.JsonSchemaFactory;
import com.networknt.schema.SpecVersion;
import com.networknt.schema.ValidationMessage;

import io.modelcontextprotocol.util.Assert;

/**
 * Default implementation of the {@link JsonSchemaValidator} interface. This class
 * provides methods to validate structured content against a JSON schema. It uses the
 * NetworkNT JSON Schema Validator library for validation.
 *
 * @author Christian Tzolov
 */
public class DefaultJsonSchemaValidator implements JsonSchemaValidator {

	private static final Logger logger = LoggerFactory.getLogger(DefaultJsonSchemaValidator.class);

	private final ObjectMapper objectMapper;

	private final JsonSchemaFactory schemaFactory;

	// TODO: Implement a strategy to purge the cache (TTL, size limit, etc.)
	private final ConcurrentHashMap<String, JsonSchema> schemaCache;

	public DefaultJsonSchemaValidator() {
		this(new ObjectMapper());
	}

	public DefaultJsonSchemaValidator(ObjectMapper objectMapper) {
		this.objectMapper = objectMapper;
		this.schemaFactory = JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012);
		this.schemaCache = new ConcurrentHashMap<>();
	}

	@Override
	public ValidationResponse validate(Map<String, Object> schema, Object structuredContent) {

		Assert.notNull(schema, "Schema must not be null");
		Assert.notNull(structuredContent, "Structured content must not be null");

		try {

			JsonNode jsonStructuredOutput = this.objectMapper.valueToTree(structuredContent);

			Set<ValidationMessage> validationResult = this.getOrCreateJsonSchema(schema).validate(jsonStructuredOutput);

			// Check if validation passed
			if (!validationResult.isEmpty()) {
				return ValidationResponse
					.asInvalid("Validation failed: structuredContent does not match tool outputSchema. "
							+ "Validation errors: " + validationResult);
			}

			return ValidationResponse.asValid(jsonStructuredOutput.toString());

		}
		catch (JsonProcessingException e) {
			logger.error("Failed to validate CallToolResult: Error parsing schema: {}", e);
			return ValidationResponse.asInvalid("Error parsing tool JSON Schema: " + e.getMessage());
		}
		catch (Exception e) {
			logger.error("Failed to validate CallToolResult: Unexpected error: {}", e);
			return ValidationResponse.asInvalid("Unexpected validation error: " + e.getMessage());
		}
	}

	/**
	 * Gets a cached JsonSchema or creates and caches a new one.
	 * @param schema the schema map to convert
	 * @return the compiled JsonSchema
	 * @throws JsonProcessingException if schema processing fails
	 */
	private JsonSchema getOrCreateJsonSchema(Map<String, Object> schema) throws JsonProcessingException {
		// Generate cache key based on schema content
		String cacheKey = this.generateCacheKey(schema);

		// Try to get from cache first
		JsonSchema cachedSchema = this.schemaCache.get(cacheKey);
		if (cachedSchema != null) {
			return cachedSchema;
		}

		// Create new schema if not in cache
		JsonSchema newSchema = this.createJsonSchema(schema);

		// Cache the schema
		JsonSchema existingSchema = this.schemaCache.putIfAbsent(cacheKey, newSchema);
		return existingSchema != null ? existingSchema : newSchema;
	}

	/**
	 * Creates a new JsonSchema from the given schema map.
	 * @param schema the schema map
	 * @return the compiled JsonSchema
	 * @throws JsonProcessingException if schema processing fails
	 */
	private JsonSchema createJsonSchema(Map<String, Object> schema) throws JsonProcessingException {
		// Convert schema map directly to JsonNode (more efficient than string
		// serialization)
		JsonNode schemaNode = this.objectMapper.valueToTree(schema);

		// Handle case where ObjectMapper might return null (e.g., in mocked scenarios)
		if (schemaNode == null) {
			throw new JsonProcessingException("Failed to convert schema to JsonNode") {
			};
		}

		// Handle additionalProperties setting
		if (schemaNode.isObject()) {
			ObjectNode objectSchemaNode = (ObjectNode) schemaNode;
			if (!objectSchemaNode.has("additionalProperties")) {
				// Clone the node before modification to avoid mutating the original
				objectSchemaNode = objectSchemaNode.deepCopy();
				objectSchemaNode.put("additionalProperties", false);
				schemaNode = objectSchemaNode;
			}
		}

		return this.schemaFactory.getSchema(schemaNode);
	}

	/**
	 * Generates a cache key for the given schema map.
	 * @param schema the schema map
	 * @return a cache key string
	 */
	protected String generateCacheKey(Map<String, Object> schema) {
		if (schema.containsKey("$id")) {
			// Use the (optional) "$id" field as the cache key if present
			return "" + schema.get("$id");
		}
		// Fall back to schema's hash code as a simple cache key
		// For more sophisticated caching, could use content-based hashing
		return String.valueOf(schema.hashCode());
	}

	/**
	 * Clears the schema cache. Useful for testing or memory management.
	 */
	public void clearCache() {
		this.schemaCache.clear();
	}

	/**
	 * Returns the current size of the schema cache.
	 * @return the number of cached schemas
	 */
	public int getCacheSize() {
		return this.schemaCache.size();
	}

}