diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/JsonDefinition.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/JsonDefinition.java index ee40ee64c8..39608393e1 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/JsonDefinition.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/JsonDefinition.java @@ -1,5 +1,6 @@ package io.stargate.sgv2.jsonapi.api.model.command; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.databind.JsonNode; import com.google.common.annotations.VisibleForTesting; import java.util.Objects; @@ -10,13 +11,15 @@ * intermediate form ({@link JsonNode}) from which actual deserialization and validation can be * deferred until we have context we need * + *

Fields are {@code @JsonIgnore} so OpenAPI and MCP describe only the clause JSON. + * * @see io.stargate.sgv2.jsonapi.api.model.command.clause.filter.FilterDefinition * @see io.stargate.sgv2.jsonapi.api.model.command.clause.filter.SortDefinition * @param Type of the value object to be deserialized from the JSON value */ public abstract class JsonDefinition { /** The wrapped JSON value */ - private final JsonNode json; + @JsonIgnore private final JsonNode json; protected JsonDefinition(JsonNode json) { this.json = Objects.requireNonNull(json); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/FilterDefinition.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/FilterDefinition.java index 549173565d..63f1157721 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/FilterDefinition.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/FilterDefinition.java @@ -1,6 +1,7 @@ package io.stargate.sgv2.jsonapi.api.model.command.clause.filter; import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.databind.JsonNode; import io.stargate.sgv2.jsonapi.api.model.command.CommandContext; import io.stargate.sgv2.jsonapi.api.model.command.JsonDefinition; @@ -14,18 +15,26 @@ * allows for lazy deserialization into a {@link FilterClause}. */ @Schema( + name = "FilterDefinition", type = SchemaType.OBJECT, - implementation = FilterClause.class, - examples = - """ - {"name": "Aaron", "country": {"$eq": "NZ"}, "age": {"$gt": 40}} - """) + // FilterClause describes parsed state and would incorrectly reserve fields such as empty. + implementation = Object.class, + nullable = true, + description = "A filter to select documents, see examples.", + examples = { + """ + {"age": {"$gt": 40}} + """, + """ + {"name": "Aaron", "country": {"$eq": "NZ"}, "age": {"$gt": 40}} + """ + }) public class FilterDefinition extends JsonDefinition { /** * Lazily deserialized {@link FilterClause} from the JSON value. We need this due to existing * reliance on specific stateful instances of {@link FilterClause}. */ - private FilterClause filterClause; + @JsonIgnore private FilterClause filterClause; /** * To deserialize the whole JSON value, need to ensure DELEGATING mode (instead of PROPERTIES). diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/SortDefinition.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/SortDefinition.java index 7b3ed5b177..88d21536fa 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/SortDefinition.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/filter/SortDefinition.java @@ -1,6 +1,7 @@ package io.stargate.sgv2.jsonapi.api.model.command.clause.filter; import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.JsonNodeFactory; import io.stargate.sgv2.jsonapi.api.model.command.CommandContext; @@ -17,18 +18,29 @@ * for lazy deserialization into a {@link SortClause}. */ @Schema( + name = "SortDefinition", type = SchemaType.OBJECT, - implementation = SortClause.class, - examples = - """ - {"user.age" : -1, "user.name" : 1} - """) + implementation = Object.class, + nullable = true, + description = + "A sort to order selected documents. Supports sorting document fields by ascending or descending, or sorting by $vector, $vectorize, or $lexical. See examples.", + examples = { + """ + {"user.age" : -1, "user.name" : 1} + """, + """ + {"$vector": [0.1, 0.2, 0.3]} + """, + """ + {"$vectorize": "search text"} + """ + }) public class SortDefinition extends JsonDefinition { /** * Lazily deserialized {@link SortClause} from the JSON value. We need this due to existing * reliance on specific stateful instances of {@link SortClause}. */ - private SortClause sortClause; + @JsonIgnore private SortClause sortClause; /** * To deserialize the whole JSON value, need to ensure DELEGATING mode (instead of PROPERTIES). @@ -84,6 +96,7 @@ public boolean hasVsearchClause() { * Helper method needed for logging purposes, where caller does not have access to the {@link * CommandContext}. */ + @JsonIgnore public List getSortExpressionPaths() { // We will either be wrapping the clause or have a json value: if (sortClause != null) { diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/sort/FindAndRerankSort.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/sort/FindAndRerankSort.java index c0f02d31eb..efc1c8fd25 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/sort/FindAndRerankSort.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/sort/FindAndRerankSort.java @@ -1,11 +1,15 @@ package io.stargate.sgv2.jsonapi.api.model.command.clause.sort; +import static io.stargate.sgv2.jsonapi.config.constants.DocumentConstants.Fields.*; + +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import io.stargate.sgv2.jsonapi.metrics.CommandFeatures; import io.stargate.sgv2.jsonapi.util.recordable.Recordable; import java.util.*; import org.eclipse.microprofile.openapi.annotations.enums.SchemaType; import org.eclipse.microprofile.openapi.annotations.media.Schema; +import org.eclipse.microprofile.openapi.annotations.media.SchemaProperty; /** * @param vectorizeSort If null then the $vectorize field was not present in the sort @@ -19,18 +23,84 @@ */ @JsonDeserialize(using = FindAndRerankSortClauseDeserializer.class) @Schema( + name = "FindAndRerankSort", type = SchemaType.OBJECT, implementation = Map.class, - examples = - """ - {"$sort" : {"$hybrid" : "Same query for vectorize and bm25 sorting"}}} - {"$sort" : {"$hybrid" : {"$vectorize" : "vectorize sort query" , "$lexical": "lexical sort" }}} - {"$sort" : {"$hybrid" : {"$vector" : [1,2,3] , "$lexical": "lexical sort" }}} - """) + nullable = true, + additionalProperties = Schema.False.class, + properties = + @SchemaProperty( + name = HYBRID_FIELD, + anyOf = {String.class, FindAndRerankSort.HybridSortSchema.class}), + description = + "Use $hybrid with a query string or an object containing $vectorize or $vector, and optionally $lexical." + + " $vector and $vectorize cannot be used together.", + examples = { + """ + {"$hybrid": "Same query for vectorize and BM25 sorting"} + """, + """ + {"$hybrid": {"$vectorize": "vectorize sort query", "$lexical": "lexical sort"}} + """, + """ + {"$hybrid": {"$vector": [1, 2, 3], "$lexical": "lexical sort"}} + """, + """ + {"$hybrid": {"$vector": {"$binary": "P4AAAEAAAABAQAAA"}, "$lexical": "lexical query"}} + """ + }) public record FindAndRerankSort( - String vectorizeSort, String lexicalSort, float[] vectorSort, CommandFeatures commandFeatures) + @JsonIgnore String vectorizeSort, + @JsonIgnore String lexicalSort, + @JsonIgnore float[] vectorSort, + @JsonIgnore CommandFeatures commandFeatures) implements Recordable { + // Schema-only classes. @SchemaProperty can point at a nested schema only through a class, so + // these empty classes carry the @Schema for each nested part of the sort. Nothing creates them. + + /** The object form of $hybrid, with $vector or $vectorize and an optional $lexical. */ + @Schema( + name = "HybridSort", + type = SchemaType.OBJECT, + implementation = Object.class, + additionalProperties = Schema.False.class, + properties = { + @SchemaProperty( + name = VECTOR_EMBEDDING_TEXT_FIELD, + type = SchemaType.STRING, + nullable = true), + @SchemaProperty(name = LEXICAL_CONTENT_FIELD, type = SchemaType.STRING, nullable = true), + @SchemaProperty( + name = VECTOR_EMBEDDING_FIELD, + nullable = true, + anyOf = {float[].class, BinaryVectorSchema.class}) + }, + not = ConflictingVectorSortSchema.class) + public static class HybridSortSchema {} + + /** A binary encoded vector such as {"$binary": "P4AAAEAAAABAQAAA"}. */ + @Schema( + name = "BinaryVector", + type = SchemaType.OBJECT, + implementation = Object.class, + nullable = true, + additionalProperties = Schema.False.class, + requiredProperties = BINARY_VECTOR_TEXT_FIELD, + properties = @SchemaProperty(name = BINARY_VECTOR_TEXT_FIELD, type = SchemaType.STRING)) + public static class BinaryVectorSchema {} + + /** + * Matches an object that has both $vector and $vectorize. HybridSort uses it as its {@code not} + * schema, so a sort that sets both fails validation. + */ + @Schema( + name = "ConflictingVectorSort", + type = SchemaType.OBJECT, + implementation = Object.class, + requiredProperties = {VECTOR_EMBEDDING_FIELD, VECTOR_EMBEDDING_TEXT_FIELD}) + public static class ConflictingVectorSortSchema {} + public static final FindAndRerankSort NO_ARG_SORT = new FindAndRerankSort(null, null, null, CommandFeatures.EMPTY); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/FindAndRerankCommand.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/FindAndRerankCommand.java index ae300a8d7d..0cab88a25e 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/FindAndRerankCommand.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/FindAndRerankCommand.java @@ -1,7 +1,13 @@ package io.stargate.sgv2.jsonapi.api.model.command.impl; import static io.stargate.sgv2.jsonapi.config.constants.DocumentConstants.Fields.*; +import static io.stargate.sgv2.jsonapi.config.constants.OpenApiConstants.Components.SCHEMA_REF_PREFIX; +import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.AUTHENTICATION; +import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.MODEL_NAME; +import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.PARAMETERS; +import static io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants.PROVIDER; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonTypeName; import com.fasterxml.jackson.core.JsonParser; @@ -28,10 +34,9 @@ import javax.annotation.Nullable; import org.eclipse.microprofile.openapi.annotations.enums.SchemaType; import org.eclipse.microprofile.openapi.annotations.media.Schema; +import org.eclipse.microprofile.openapi.annotations.media.SchemaProperty; -@Schema( - description = - "Finds documents using using vector and lexical sorting, then reranks the results.") +@Schema(description = "Finds documents using vector and lexical sorting, then reranks the results.") @JsonTypeName(CommandName.Names.FIND_AND_RERANK) public record FindAndRerankCommand( @Valid @JsonProperty("filter") FilterDefinition filterDefinition, @@ -70,15 +75,7 @@ public record Options( implementation = Integer.class) Integer limit, /** ---- */ - @Schema( - description = - "The maximum number of documents to read for the vector and lexical queries that feed into the reranking. May be a number or an object with $vector and $lexical fields. The accepted range is determined by server configuration.", - examples = - """ - {"hybridLimits" : 100} - {"hybridLimits" : {"$vector" : 100, "$lexical" : 10}} - """) - HybridLimits hybridLimits, + @Schema(ref = HybridLimits.REF) HybridLimits hybridLimits, /** ---- */ @Schema( description = @@ -117,18 +114,67 @@ public record Options( "Optional override for the reranking service. Completely replaces the" + " collection-level reranking configuration when provided. Both 'provider' and" + " 'modelName' are required.", - implementation = CreateCollectionCommand.Options.RerankServiceDesc.class) + implementation = RerankOverrideSchema.class, + ref = RerankOverrideSchema.REF) @JsonProperty("rerank") CreateCollectionCommand.Options.RerankServiceDesc rerankServiceOverride) {} + // Matches the override validation in CollectionRerankDef. Update both together. + @Schema( + name = RerankOverrideSchema.NAME, + type = SchemaType.OBJECT, + implementation = Object.class, + nullable = true, + additionalProperties = Schema.False.class, + requiredProperties = {PROVIDER, MODEL_NAME}, + properties = { + @SchemaProperty(name = PROVIDER, type = SchemaType.STRING), + @SchemaProperty(name = MODEL_NAME, type = SchemaType.STRING), + @SchemaProperty( + name = AUTHENTICATION, + type = SchemaType.OBJECT, + nullable = true, + additionalProperties = String.class), + @SchemaProperty( + name = PARAMETERS, + type = SchemaType.OBJECT, + nullable = true, + additionalProperties = Schema.False.class) + }, + examples = + """ + {"provider": "nvidia", "modelName": "nvidia/llama-3.2-nv-rerankqa-1b-v2"} + """) + public static class RerankOverrideSchema { + public static final String NAME = "RerankServiceOverride"; + public static final String REF = SCHEMA_REF_PREFIX + NAME; + } + @JsonDeserialize(using = HybridLimitsDeserializer.class) + @Schema( + name = HybridLimits.NAME, + description = + "Candidate read limits as one integer for both legs, or an object containing both $vector and $lexical integer limits.", + // Void is the annotation default, but the OpenAPI scanner only sees values set in source. + // Setting it stops the scanner adding the record fields as an object type next to anyOf. + implementation = Void.class, + anyOf = {Integer.class, HybridLimits.ValuesSchema.class}, + examples = { + "50", + """ + {"$vector": 50, "$lexical": 10} + """ + }) public record HybridLimits( @JsonProperty(DocumentConstants.Fields.VECTOR_EMBEDDING_FIELD) int vectorLimit, /** ---- */ @JsonProperty(DocumentConstants.Fields.LEXICAL_CONTENT_FIELD) int lexicalLimit, - CommandFeatures commandFeatures) + @JsonIgnore CommandFeatures commandFeatures) implements Recordable { + public static final String NAME = "HybridLimits"; + public static final String REF = SCHEMA_REF_PREFIX + NAME; + public static final HybridLimits DEFAULT = new HybridLimits( RerankingConstants.HybridSearchLimits.DEFAULT, @@ -142,6 +188,19 @@ public DataRecorder recordTo(DataRecorder dataRecorder) { .append("lexicalLimit", lexicalLimit) .append("commandFeatures", commandFeatures); } + + /** Schema for the object form of the limits, the second branch of the anyOf above. */ + @Schema( + name = "HybridLimitValues", + type = SchemaType.OBJECT, + nullable = true, + additionalProperties = Schema.False.class, + requiredProperties = {VECTOR_EMBEDDING_FIELD, LEXICAL_CONTENT_FIELD}, + properties = { + @SchemaProperty(name = VECTOR_EMBEDDING_FIELD, type = SchemaType.INTEGER), + @SchemaProperty(name = LEXICAL_CONTENT_FIELD, type = SchemaType.INTEGER) + }) + public static class ValuesSchema {} } /** Deserializer for the `hybridLimits` option, the limits for the inner reads. */ diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/CollectionCommandTools.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/CollectionCommandTools.java index 529fc910ec..efb89bf227 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/CollectionCommandTools.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/CollectionCommandTools.java @@ -148,16 +148,23 @@ public Uni estimatedDocumentCount( return mcpResource.processCollectionCommand(keyspace, collection, command); } - @Tool( - description = - "Finds documents using using vector and lexical sorting, then reranks the results.") + @Tool(description = "Finds documents using vector and lexical sorting, then reranks the results.") public Uni findAndRerank( @ToolArg(description = "Name of the keyspace") String keyspace, - @ToolArg(description = "Name of the collection/table") String collection, - @ToolArg(description = "filter", required = false) FilterDefinition filter, - @ToolArg(description = "projection", required = false) JsonNode projection, - @ToolArg(description = "sort", required = false) FindAndRerankSort sort, - @ToolArg(description = "options", required = false) FindAndRerankCommand.Options options) { + @ToolArg(description = "Name of the collection") String collection, + // filter and sort have no description here because it would replace the one on their + // classes, which Swagger also uses. + @ToolArg(required = false) FilterDefinition filter, + @ToolArg( + description = + "Fields to include or exclude in the returned documents, for example {\"title\": 1}", + required = false) + JsonNode projection, + @ToolArg(required = false) FindAndRerankSort sort, + @ToolArg( + description = "Options for result limits, reranking, and what the response includes", + required = false) + FindAndRerankCommand.Options options) { var command = new FindAndRerankCommand(filter, projection, sort, options); return mcpResource.processCollectionCommand(keyspace, collection, command); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpOpenApiSchemaAdapter.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpOpenApiSchemaAdapter.java new file mode 100644 index 0000000000..95207f1039 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpOpenApiSchemaAdapter.java @@ -0,0 +1,304 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static io.stargate.sgv2.jsonapi.config.constants.OpenApiConstants.Components.SCHEMAS_POINTER; +import static io.stargate.sgv2.jsonapi.config.constants.OpenApiConstants.Components.SCHEMA_REF_PREFIX; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.quarkiverse.mcp.server.GlobalInputSchemaGenerator; +import io.quarkiverse.mcp.server.ToolManager.ToolInfo; +import io.quarkus.smallrye.openapi.runtime.OpenApiConstants; +import io.vertx.core.json.JsonObject; +import jakarta.annotation.Priority; +import jakarta.decorator.Decorator; +import jakarta.decorator.Delegate; +import jakarta.inject.Inject; +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; +import java.util.function.Consumer; + +/** + * Makes each MCP tool input schema standalone by copying in the OpenAPI schemas it points to. + * + *

{@link McpSchemaCustomizer} replaces the schema of each clause argument, such as filter or + * sort, with a pointer to the schema Quarkus generates for Swagger, so both APIs share one + * definition. An MCP tool's input schema is a standalone JSON Schema document, though, so a pointer + * into the OpenAPI document leads nowhere for an MCP client. + * + *

After the MCP library generates a tool schema, this class copies each OpenAPI schema the tool + * needs into the tool schema's {@code $defs}, and rewrites the pointers to match: + * + * 

+ * before: {"properties": {"filter": {"$ref": "#/components/schemas/FilterDefinition"}}}
+ * after:  {"properties": {"filter": {"$ref": "#/$defs/_openapi_FilterDefinition"}},
+ *          "$defs": {"_openapi_FilterDefinition": (copy of the OpenAPI schema)}}
+ *
+ * + *

Pointers stay pointers rather than being replaced by the schemas they name. Clause schemas can + * contain themselves, so replacing them would never finish, and keywords written next to a pointer + * would have to be merged into the copied schema. + */ +// @Decorator wraps the MCP library's generator instead of replacing it. CDI enables a decorator for +// the whole application only when it has @Priority. The value orders decorators of the same bean, +// and this is the only one. +@Decorator +@Priority(1) +public class McpOpenApiSchemaAdapter implements GlobalInputSchemaGenerator { + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final TypeReference> JSON_OBJECT = new TypeReference<>() {}; + + /** Start of a pointer to a definition in the tool schema's own {@code $defs}. */ + private static final String DEFS_REF_PREFIX = "#/$defs/"; + + /** + * Start of each copied schema's name in {@code $defs}. The MCP library names its own definitions + * after Java types, so the prefix keeps a copied {@code FilterDefinition} apart from a definition + * the library generated from a Java class with the same name. + */ + private static final String DEFINITION_PREFIX = "_openapi_"; + + // The JSON Schema 2020-12 keywords that hold other schemas. The walk follows only these, so a + // "$ref" that appears as data, in an example or a default value, is never rewritten. + + /** Keywords whose value is an object of schemas, such as {@code properties}. */ + private static final Set SCHEMA_OBJECTS = + Set.of("properties", "patternProperties", "$defs", "dependentSchemas"); + + /** Keywords whose value is an array of schemas, such as {@code anyOf}. */ + private static final Set SCHEMA_ARRAYS = Set.of("allOf", "anyOf", "oneOf", "prefixItems"); + + /** Keywords whose value is a single schema, such as {@code items}. */ + private static final Set SCHEMA_VALUES = + Set.of( + "items", + "additionalProperties", + "unevaluatedProperties", + "contains", + "not", + "propertyNames", + "if", + "then", + "else", + "unevaluatedItems", + "contentSchema"); + + private final GlobalInputSchemaGenerator delegate; + + /** The OpenAPI document Quarkus generated at build time. Read once and never changed. */ + private final JsonNode openApi; + + /** + * Reads the OpenAPI document once, here, rather than on each request. The MCP HTTP extension + * creates its tool manager, and with it this decorator, while the server starts. A missing + * document therefore stops startup instead of failing each request that lists tools. + */ + @Inject + public McpOpenApiSchemaAdapter(@Delegate GlobalInputSchemaGenerator delegate) { + this.delegate = delegate; + this.openApi = readGeneratedOpenApiDocument(); + } + + /** + * Generates the tool schema as usual, then copies in the OpenAPI schemas it points to. + * + *

The MCP library calls this each time a client lists tools, once for each tool. + */ + @Override + public InputSchema generate(ToolInfo tool) { + ObjectNode schema; + try { + // The interface promises the schema only as JSON text, so read it once to change it. + schema = (ObjectNode) MAPPER.readTree(delegate.generate(tool).asJson()); + } catch (JsonProcessingException e) { + // Not expected: the library wrote this JSON itself. Listing tools fails with this error. + throw new IllegalStateException("Cannot read the MCP input schema for " + tool.name(), e); + } + bundleOpenApiSchemas(schema, openApi); + // The library sends value(), a Vert.x object. Convert the tree to it without writing JSON text. + return new StandaloneInputSchema(new JsonObject(MAPPER.convertValue(schema, JSON_OBJECT))); + } + + /** + * Changes {@code schema} in place so it no longer points into the OpenAPI document. + * + *

Package-private and static so tests can check the conversion against a small OpenAPI + * document without starting Quarkus. {@link #generate} is the only production caller. + * + * @param schema a tool input schema, which this method changes + * @param openApi the OpenAPI document, which this method only reads + * @throws IllegalStateException if a pointer names a schema the OpenAPI document does not have + */ + static void bundleOpenApiSchemas(ObjectNode schema, JsonNode openApi) { + rebaseArgumentReferences(schema); + var copies = copyReferencedSchemas(schema, openApi); + if (copies.isEmpty()) { + // Most tools have no clause arguments. Leave their schemas without an empty $defs. + return; + } + var definitions = schema.withObject("/$defs"); + copies.forEach( + (name, copy) -> { + // Library definitions never start with the prefix, so a clash is a bug in this class. + // Fail rather than replace a definition. Listing tools fails with this error. + if (definitions.has(name)) { + throw new IllegalStateException("Duplicate MCP schema definition " + name); + } + definitions.set(name, copy); + }); + } + + /** Returns a pointer to the named schema under {@code components/schemas} in OpenAPI. */ + static String componentRef(String name) { + return SCHEMA_REF_PREFIX + escapePointerToken(name); + } + + /** + * Reads the OpenAPI document that Quarkus generates at build time and packages with the + * application. Quarkus builds the document Swagger serves from the same file. + * + *

The lookup matches Quarkus's own {@code OpenApiDocumentService}: the class loader the + * OpenAPI extension records while starting, or the current thread's class loader if none is + * recorded. Package-private so tests can use the same packaged document. + */ + static JsonNode readGeneratedOpenApiDocument() { + var loader = OpenApiConstants.classLoader; + if (loader == null) { + loader = Thread.currentThread().getContextClassLoader(); + } + try (var stream = loader.getResourceAsStream(OpenApiConstants.BASE_NAME + ".JSON")) { + if (stream == null) { + // Only if the build stops generating the document, for example when OpenAPI generation is + // turned off. Clause schemas cannot be described without it, so stop startup. + throw new IllegalStateException("The generated OpenAPI document is unavailable"); + } + return MAPPER.readTree(stream); + } catch (IOException e) { + throw new IllegalStateException("Cannot read the generated OpenAPI document", e); + } + } + + /** + * Fixes pointers inside tool arguments that carry their own {@code $defs}. + * + *

The MCP library generates each tool argument separately and nests the result under {@code + * properties}. When an argument's schema has its own {@code $defs}, its pointers still start at + * the document root, where those definitions do not exist: + * + * 

+   * before: {"properties": {"options": {"$defs": {"Local": ...},
+   *                                     "properties": {"value": {"$ref": "#/$defs/Local"}}}}}
+   * after:  the inner pointer becomes "#/properties/options/$defs/Local"
+   *
+ * + *

A pointer changes only when the root lacks the definition and the argument has it, so + * pointers that already work stay as they are. Each argument keeps its own definitions, so two + * arguments can both define {@code Local} without a clash. + */ + private static void rebaseArgumentReferences(ObjectNode schema) { + for (var argument : schema.path("properties").properties()) { + var argumentSchema = argument.getValue(); + if (!argumentSchema.path("$defs").isObject()) { + continue; + } + var argumentPointer = "#/properties/" + escapePointerToken(argument.getKey()); + visitSchemas( + argumentSchema, + node -> { + var reference = node.path("$ref").asText(); + if (reference.startsWith(DEFS_REF_PREFIX) + && schema.at(reference.substring(1)).isMissingNode() + && !argumentSchema.at(reference.substring(1)).isMissingNode()) { + node.put("$ref", argumentPointer + reference.substring(1)); + } + }); + } + } + + /** + * Rewrites each OpenAPI pointer in {@code schema} to point into {@code $defs}. Returns copies of + * the OpenAPI schemas those pointers need, keyed by their {@code $defs} name. + * + *

Copied schemas can point to further OpenAPI schemas, so they are rewritten and searched as + * well. Each schema is copied once, which also ends the search for schemas that contain + * themselves. + */ + private static Map copyReferencedSchemas(ObjectNode schema, JsonNode openApi) { + var copies = new LinkedHashMap(); + var pending = new ArrayDeque(); + pending.add(schema); + while (!pending.isEmpty()) { + visitSchemas( + pending.remove(), + node -> { + var reference = node.path("$ref").asText(); + if (!reference.startsWith(SCHEMA_REF_PREFIX)) { + return; + } + if (openApi.at(reference.substring(1)).isMissingNode()) { + throw new IllegalStateException("Unresolved OpenAPI schema reference " + reference); + } + // The path is usually a schema name, such as "FilterDefinition", but it can point + // inside one, such as "FilterDefinition/properties/x". Copy the whole named schema and + // keep the rest of the path. + var path = reference.substring(SCHEMA_REF_PREFIX.length()); + var escapedName = path.split("/", 2)[0]; + node.put("$ref", DEFS_REF_PREFIX + DEFINITION_PREFIX + path); + // Pointers escape "~" and "/" in names, but $defs keys hold the plain name. + var name = DEFINITION_PREFIX + unescapePointerToken(escapedName); + if (!copies.containsKey(name)) { + // Copy, because every request shares the document and the copy gets rewritten. + var copy = openApi.at(SCHEMAS_POINTER + escapedName).deepCopy(); + copies.put(name, copy); + pending.add(copy); + } + }); + } + return copies; + } + + /** Calls {@code visitor} on {@code node} and on every schema inside it, parents first. */ + private static void visitSchemas(JsonNode node, Consumer visitor) { + if (!(node instanceof ObjectNode schema)) { + return; + } + visitor.accept(schema); + for (var keyword : SCHEMA_OBJECTS) { + if (schema.path(keyword).isObject()) { + schema.path(keyword).forEach(value -> visitSchemas(value, visitor)); + } + } + for (var keyword : SCHEMA_ARRAYS) { + if (schema.path(keyword).isArray()) { + schema.path(keyword).forEach(value -> visitSchemas(value, visitor)); + } + } + for (var keyword : SCHEMA_VALUES) { + visitSchemas(schema.path(keyword), visitor); + } + } + + /** Escapes a name for a JSON Pointer (RFC 6901): "~" becomes "~0" and "/" becomes "~1". */ + private static String escapePointerToken(String name) { + return name.replace("~", "~0").replace("/", "~1"); + } + + /** Reverses {@link #escapePointerToken}, replacing "~1" first as RFC 6901 requires. */ + private static String unescapePointerToken(String token) { + return token.replace("~1", "/").replace("~0", "~"); + } + + /** The tool schema as the Vert.x object the MCP library puts into its tool list. */ + private record StandaloneInputSchema(JsonObject value) implements InputSchema { + @Override + public String asJson() { + // Matches the library's own InputSchema implementation. + return value.toString(); + } + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaCustomizer.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaCustomizer.java new file mode 100644 index 0000000000..5dbe9c9534 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaCustomizer.java @@ -0,0 +1,107 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.node.JsonNodeFactory; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.github.victools.jsonschema.generator.CustomDefinition; +import com.github.victools.jsonschema.generator.CustomPropertyDefinition; +import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder; +import io.quarkiverse.mcp.server.runtime.SchemaGeneratorConfigCustomizer; +import jakarta.inject.Singleton; +import org.eclipse.microprofile.openapi.annotations.media.Schema; + +/** + * Generates MCP schemas using Data API clause definitions and Jackson and OpenAPI annotations. + * + *

The MCP server builds its generator with the victools PLAIN_JSON preset and no modules, so by + * default it reads no annotations at all. Without this class, MCP schemas would use Java field + * names, include fields Jackson ignores, and drop every {@code @Schema} description. + * + *

MCP clients use these schemas to build tool arguments. Custom deserialization can accept JSON + * shapes that differ from the Java fields, so those types need overrides. Ordinary records still + * use field-based generation. + * + *

Named {@code @Schema} definitions refer to the schemas Quarkus already generates for Swagger. + * {@link McpOpenApiSchemaAdapter} packages those definitions into the complete MCP input schema. + * + *

This class also includes the field description resolver that was previously in {@code + * McpSchemaDescriptionCustomizer}. + */ +@Singleton +public class McpSchemaCustomizer implements SchemaGeneratorConfigCustomizer { + @Override + public void customize(SchemaGeneratorConfigBuilder builder) { + builder + .forFields() + // Field descriptions come from @Schema, as they do for Swagger. + .withDescriptionResolver( + field -> { + var schema = field.getAnnotationConsideringFieldAndGetter(Schema.class); + return schema != null && !schema.description().isBlank() + ? schema.description() + : null; + }) + // Leave out fields that are not part of the request JSON. + .withIgnoreCheck( + field -> { + var ignore = field.getAnnotationConsideringFieldAndGetter(JsonIgnore.class); + var schema = field.getAnnotationConsideringFieldAndGetter(Schema.class); + return (ignore != null && ignore.value()) || (schema != null && schema.hidden()); + }) + // A field with @Schema(ref) uses that OpenAPI component instead of its Java type. + .withCustomDefinitionProvider( + (field, context) -> { + var schema = + schemaNode(field.getAnnotationConsideringFieldAndGetter(Schema.class), false); + return schema == null + ? null + : new CustomPropertyDefinition(schema, CustomDefinition.AttributeInclusion.YES); + }) + // Property names come from @JsonProperty, so MCP sees the JSON names clients send. + .withPropertyNameOverrideResolver( + field -> { + var property = field.getAnnotationConsideringFieldAndGetter(JsonProperty.class); + // Jackson uses an empty value for the Java name. Whitespace is an explicit JSON name. + return property != null && !property.value().isEmpty() ? property.value() : null; + }); + + builder + .forTypesInGeneral() + // Type descriptions come from @Schema on the class. + .withDescriptionResolver( + type -> { + var schema = type.getType().getErasedType().getAnnotation(Schema.class); + return schema != null && !schema.description().isBlank() + ? schema.description() + : null; + }) + // A type with a named @Schema uses that OpenAPI component instead of its Java fields. + .withCustomDefinitionProvider( + (javaType, context) -> { + var schema = schemaNode(javaType.getErasedType().getAnnotation(Schema.class), true); + return schema == null + ? null + : new CustomDefinition( + schema, + CustomDefinition.DefinitionType.INLINE, + CustomDefinition.AttributeInclusion.YES); + }); + } + + private static ObjectNode schemaNode(Schema annotation, boolean typeAnnotation) { + if (annotation == null) { + return null; + } + var reference = annotation.ref(); + if (reference.isBlank()) { + // A type name identifies an OpenAPI component. A field name only names that property. + var name = typeAnnotation ? annotation.name() : ""; + if (name.isBlank()) { + return null; + } + reference = McpOpenApiSchemaAdapter.componentRef(name); + } + return JsonNodeFactory.instance.objectNode().put("$ref", reference); + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaDescriptionCustomizer.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaDescriptionCustomizer.java deleted file mode 100644 index 6e5e66736d..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaDescriptionCustomizer.java +++ /dev/null @@ -1,33 +0,0 @@ -package io.stargate.sgv2.jsonapi.api.v1.mcp; - -import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder; -import io.quarkiverse.mcp.server.runtime.SchemaGeneratorConfigCustomizer; -import jakarta.inject.Singleton; -import org.eclipse.microprofile.openapi.annotations.media.Schema; - -/** - * Enables MicroProfile {@link Schema#description()} resolution for the victools JSON schema - * generator used by Quarkus MCP Server. Without this customizer, nested record fields annotated - * with {@code @Schema(description = "...")} will not have their descriptions included in the MCP - * tool schema output. - * - *

The Quarkus MCP Server's {@code DefaultSchemaGenerator} uses {@code OptionPreset.PLAIN_JSON}, - * which does not read any annotations by default. This customizer bridges that gap by adding a - * field-level description resolver for MicroProfile OpenAPI's {@code @Schema} annotation. - */ -@Singleton -public class McpSchemaDescriptionCustomizer implements SchemaGeneratorConfigCustomizer { - - @Override - public void customize(SchemaGeneratorConfigBuilder builder) { - builder - .forFields() - .withDescriptionResolver( - field -> { - Schema schemaAnn = field.getAnnotationConsideringFieldAndGetter(Schema.class); - return schemaAnn != null && !schemaAnn.description().isEmpty() - ? schemaAnn.description() - : null; - }); - } -} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/constants/OpenApiConstants.java b/src/main/java/io/stargate/sgv2/jsonapi/config/constants/OpenApiConstants.java index 9412cff336..87e6fc1c5d 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/constants/OpenApiConstants.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/constants/OpenApiConstants.java @@ -8,4 +8,14 @@ interface SecuritySchemes { String TOKEN = "Token"; } + + /** Locations of the reusable schemas in the OpenAPI document. */ + interface Components { + + /** JSON Pointer to the reusable schemas, for reading them from the document. */ + String SCHEMAS_POINTER = "/components/schemas/"; + + /** Start of a {@code $ref} to a reusable schema. Add the schema name to the end. */ + String SCHEMA_REF_PREFIX = "#" + SCHEMAS_POINTER; + } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpClauseSchemaContractTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpClauseSchemaContractTest.java new file mode 100644 index 0000000000..99f5e33de0 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpClauseSchemaContractTest.java @@ -0,0 +1,686 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static io.restassured.RestAssured.given; +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.node.MissingNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.github.victools.jsonschema.generator.*; +import com.networknt.schema.JsonSchemaFactory; +import com.networknt.schema.SpecVersion; +import io.quarkiverse.mcp.server.GlobalInputSchemaGenerator; +import io.quarkiverse.mcp.server.ToolManager; +import io.quarkus.test.junit.QuarkusTest; +import io.quarkus.test.junit.TestProfile; +import io.stargate.sgv2.jsonapi.TestConstants; +import io.stargate.sgv2.jsonapi.api.model.command.JsonDefinition; +import io.stargate.sgv2.jsonapi.api.model.command.clause.filter.FilterDefinition; +import io.stargate.sgv2.jsonapi.api.model.command.clause.filter.SortDefinition; +import io.stargate.sgv2.jsonapi.api.model.command.clause.sort.FindAndRerankSort; +import io.stargate.sgv2.jsonapi.api.model.command.impl.CreateCollectionCommand; +import io.stargate.sgv2.jsonapi.api.model.command.impl.FindAndRerankCommand; +import io.stargate.sgv2.jsonapi.exception.RequestException; +import io.stargate.sgv2.jsonapi.metrics.CommandFeatures; +import io.stargate.sgv2.jsonapi.service.provider.ApiModelSupport; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig.RerankingProviderConfig.AuthenticationType; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfigImpl; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfigImpl.RerankingProviderConfigImpl; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfigImpl.RerankingProviderConfigImpl.AuthenticationConfigImpl; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfigImpl.RerankingProviderConfigImpl.ModelConfigImpl; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfigImpl.RerankingProviderConfigImpl.ModelConfigImpl.RequestPropertiesImpl; +import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionRerankDef; +import io.stargate.sgv2.jsonapi.testresource.NoGlobalResourcesTestProfile; +import jakarta.inject.Inject; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Stream; +import org.eclipse.microprofile.openapi.annotations.media.Schema; +import org.jboss.jandex.AnnotationTarget; +import org.jboss.jandex.IndexReader; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; + +/** + * Checks that MCP and OpenAPI advertise the JSON accepted by Data API clause deserializers. + * Reflection alone exposes internal Java fields for clauses with custom deserializers, so those + * clauses need schema overrides. + * + *

Schema validation checks the advertised JSON shape. Deserialization and clause building check + * that the server accepts the examples. Generated tool schemas and the OpenAPI document also verify + * that the overrides reach clients. These checks do not execute commands or require a database. + */ +@QuarkusTest +@TestProfile(NoGlobalResourcesTestProfile.Impl.class) +class McpClauseSchemaContractTest { + @Inject ObjectMapper objectMapper; + @Inject GlobalInputSchemaGenerator schemaGenerator; + @Inject ToolManager toolManager; + + record WireNames( + @JsonProperty("wireName") String internalName, @JsonIgnore CommandFeatures commandFeatures) {} + + record PropertyNameDefaults( + @JsonProperty("") String defaultName, @JsonProperty(" ") String whitespaceName) {} + + record IgnoredFields( + @JsonIgnore String internal, + @JsonIgnore(false) String visible, + CommandFeatures unannotatedFeatures) {} + + static class GetterIgnoredField { + public String internal; + + @JsonIgnore + public String getInternal() { + return internal; + } + } + + @Test + void jsonIgnoreControlsFieldVisibility() { + var properties = schemaFor(IgnoredFields.class).path("properties"); + assertFalse(properties.has("internal")); + assertTrue(properties.has("visible")); + assertTrue(properties.has("unannotatedFeatures")); + } + + @Test + void jsonIgnoreOnGetterHidesField() { + assertFalse(schemaFor(GetterIgnoredField.class).path("properties").has("internal")); + } + + @Test + void fieldsUseJsonPropertyNamesAndHideCommandFeatures() { + var properties = schemaFor(WireNames.class).path("properties"); + assertTrue(properties.has("wireName")); + assertFalse(properties.has("internalName")); + assertFalse(properties.has("commandFeatures")); + } + + @Test + void emptyAndWhitespacePropertyNamesMatchJackson() throws Exception { + var payload = + objectMapper.readTree( + """ + {"defaultName": "default", " ": "explicit"} + """); + var schema = schemaFor(PropertyNameDefaults.class); + var properties = schema.path("properties"); + assertEquals(payload.size(), properties.size()); + payload.fieldNames().forEachRemaining(name -> assertTrue(properties.has(name))); + assertSchemaAccepts(schema, payload); + assertEquals( + new PropertyNameDefaults("default", "explicit"), + objectMapper.treeToValue(payload, PropertyNameDefaults.class)); + } + + /** An example must satisfy both its advertised schema and the server's clause parser. */ + @ParameterizedTest + @MethodSource("clauseTypes") + void everyAdvertisedClauseExampleDeserializes(Class clauseType) throws Exception { + var schema = schemaFor(clauseType); + var examples = schemaMetadata(schema, schema, "examples"); + assertFalse( + examples.isMissingNode(), "Clause schema must advertise examples of the accepted JSON"); + assertFalse(examples.isEmpty()); + for (var example : examples) { + assertSchemaAccepts(schema, example); + var clause = objectMapper.treeToValue(example, clauseType); + assertNotNull(clause); + // These two clauses defer parsing until build, so deserialization alone cannot validate them. + if (clause instanceof FilterDefinition filter) { + assertNotNull(filter.build(new TestConstants().collectionContext())); + } else if (clause instanceof SortDefinition sort) { + assertNotNull(sort.build(new TestConstants().collectionContext())); + } + } + } + + static Stream> clauseTypes() throws Exception { + var loader = FilterDefinition.class.getClassLoader(); + // Use the build's class index so new annotated clauses join these tests automatically. + try (var indexStream = loader.getResourceAsStream("META-INF/jandex.idx")) { + assertNotNull(indexStream, "The application class index must be available"); + var index = new IndexReader(indexStream).read(); + var clauseTypes = new ArrayList>(); + for (var annotation : index.getAnnotations(Schema.class)) { + if (annotation.target().kind() != AnnotationTarget.Kind.CLASS) { + continue; + } + var type = Class.forName(annotation.target().asClass().name().toString(), false, loader); + var schema = type.getAnnotation(Schema.class); + // Select parsed clause types, excluding helper classes for nested JSON shapes. + if (!schema.name().isBlank() + && (JsonDefinition.class.isAssignableFrom(type) + || type.isAnnotationPresent(JsonDeserialize.class))) { + clauseTypes.add(type); + } + } + assertFalse(clauseTypes.isEmpty(), "The application index must contain clause schemas"); + return clauseTypes.stream().sorted(Comparator.comparing(Class::getName)); + } + } + + /** Both entry points accept the same clause JSON and must expose the same schema definition. */ + @ParameterizedTest + @MethodSource("clauseTypes") + void openApiAndMcpShareClauseSchemas(Class clauseType) throws Exception { + var name = clauseType.getAnnotation(Schema.class).name(); + var document = openApi(); + var openApiSchema = document.path("components").path("schemas").path(name); + assertFalse(openApiSchema.isMissingNode(), "Missing OpenAPI component " + name); + var expected = + objectMapper.createObjectNode().put("$ref", McpOpenApiSchemaAdapter.componentRef(name)); + McpOpenApiSchemaAdapter.bundleOpenApiSchemas(expected, document); + // Normalize numeric examples through the same mapper before comparing schema trees. + var actual = objectMapper.readTree(schemaFor(clauseType).toString()); + assertEquals(referencedSchema(expected, expected), referencedSchema(actual, actual)); + assertEquals(expected.path("$defs"), actual.path("$defs")); + } + + @Test + void filterSchemaAcceptsDocumentObjectsAndNull() throws Exception { + var annotation = FilterDefinition.class.getAnnotation(Schema.class); + var schema = schemaFor(FilterDefinition.class); + var component = referencedSchema(schema, schema); + + assertEquals(annotation.description(), schemaMetadata(schema, schema, "description").asText()); + assertEquals(annotation.examples().length, component.path("examples").size()); + for (var i = 0; i < annotation.examples().length; i++) { + assertEquals( + objectMapper.readTree(annotation.examples()[i]), component.path("examples").get(i)); + } + assertSchemaAccepts(schema, objectMapper.readTree("null")); + assertSchemaAccepts(schema, objectMapper.readTree("{\"some.path\": {\"$gt\": 40}}")); + assertFalse( + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012) + .getSchema(schema) + .validate(objectMapper.readTree("[]")) + .isEmpty()); + } + + /** A correct reusable schema has no effect unless command fields actually reference it. */ + @Test + void openApiCommandsReferenceSharedSchemas() throws Exception { + var document = openApi(); + var schemas = document.path("components").path("schemas"); + var find = schemas.path("FindCommand").path("properties"); + assertEquals(schemas.path("FilterDefinition"), referencedSchema(document, find.path("filter"))); + assertEquals(schemas.path("SortDefinition"), referencedSchema(document, find.path("sort"))); + var rerank = schemas.path("FindAndRerankCommand").path("properties"); + assertEquals( + schemas.path("FindAndRerankSort"), referencedSchema(document, rerank.path("sort"))); + var options = referencedSchema(document, rerank.path("options")).path("properties"); + assertEquals( + schemas.path("HybridLimits"), referencedSchema(document, options.path("hybridLimits"))); + var openApiOverride = + (ObjectNode) referencedSchema(document, options.path("rerank")).deepCopy(); + // The field description belongs to the command options, not the reusable schema component. + openApiOverride.remove("description"); + assertEquals(schemas.path("RerankServiceOverride"), openApiOverride); + var mcpDocument = inputSchema("findAndRerank"); + var mcpOverride = + (ObjectNode) + referencedSchema(mcpDocument, propertySchema(mcpDocument, "options", "rerank")) + .deepCopy(); + ((ObjectNode) mcpOverride).remove("description"); + var expected = + objectMapper.createObjectNode().put("$ref", FindAndRerankCommand.RerankOverrideSchema.REF); + McpOpenApiSchemaAdapter.bundleOpenApiSchemas(expected, document); + assertEquals(referencedSchema(expected, expected), mcpOverride); + } + + /** Check the complete command schema because inferred parent types can restrict valid clauses. */ + @ParameterizedTest + @ValueSource(strings = {"50", "{\"$vector\": 50, \"$lexical\": 10}", "null"}) + void openApiAcceptsHybridLimitsWireForms(String limits) throws Exception { + var document = openApi(); + var schema = + objectMapper.createObjectNode().put("$ref", "#/components/schemas/FindAndRerankCommand"); + schema.set("components", document.path("components")); + var command = objectMapper.createObjectNode(); + command.putObject("sort").put("$hybrid", "search text"); + command.putObject("options").set("hybridLimits", objectMapper.readTree(limits)); + assertSchemaAccepts(schema, command); + } + + @Test + void openApiAcceptsNullOptionalClauses() throws Exception { + var schema = + objectMapper.createObjectNode().put("$ref", "#/components/schemas/FindAndRerankCommand"); + schema.set("components", openApi().path("components")); + assertSchemaAccepts( + schema, + objectMapper.readTree( + """ + {"filter": null, "sort": null, "options": {"rerank": null}} + """)); + } + + private JsonNode referencedSchema(JsonNode document, JsonNode schema) { + var visited = new HashSet(); + while (schema.has("$ref")) { + var reference = schema.path("$ref").asText(); + assertTrue(reference.startsWith("#/"), "Expected a local schema reference"); + assertTrue(visited.add(reference), "Circular schema reference"); + schema = document.at(reference.substring(1)); + assertFalse(schema.isMissingNode(), "Unresolved schema reference " + reference); + } + return schema; + } + + private JsonNode propertySchema(JsonNode document, String... names) { + var schema = document; + for (var name : names) { + schema = referencedSchema(document, schema).path("properties").path(name); + assertFalse(schema.isMissingNode(), "Missing property " + name); + } + return schema; + } + + private JsonNode schemaMetadata(JsonNode document, JsonNode schema, String keyword) { + if (schema.has(keyword)) { + return schema.get(keyword); + } + if (schema.has("$ref")) { + return schemaMetadata(document, referencedSchema(document, schema), keyword); + } + for (var branch : schema.path("allOf")) { + var metadata = schemaMetadata(document, branch, keyword); + if (!metadata.isMissingNode()) { + return metadata; + } + } + return MissingNode.getInstance(); + } + + private JsonNode openApi() throws Exception { + var json = + given() + .queryParam("format", "json") + .get("/api/json/openapi") + .then() + .statusCode(200) + .extract() + .asString(); + return objectMapper.readTree(json); + } + + @Test + void rerankToolUsesSortClassDescription() throws Exception { + var description = FindAndRerankSort.class.getAnnotation(Schema.class).description(); + assertFalse(description.isBlank()); + var document = inputSchema("findAndRerank"); + assertEquals( + description, + schemaMetadata(document, propertySchema(document, "sort"), "description").asText()); + } + + @Test + void rerankToolUsesFilterClassDescription() throws Exception { + var description = FilterDefinition.class.getAnnotation(Schema.class).description(); + assertFalse(description.isBlank()); + var document = inputSchema("findAndRerank"); + assertEquals( + description, + schemaMetadata(document, propertySchema(document, "filter"), "description").asText()); + } + + /** A description that only repeats the argument name tells a client nothing. */ + @Test + void rerankToolDescribesEveryArgument() throws Exception { + var document = inputSchema("findAndRerank"); + for (var argument : document.path("properties").properties()) { + var name = argument.getKey(); + var description = schemaMetadata(document, argument.getValue(), "description").asText(); + assertFalse(description.isBlank(), name); + assertNotEquals(name, description.toLowerCase(), name); + } + } + + @Test + void rerankToolUsesOverrideFieldDescription() throws Exception { + var description = + FindAndRerankCommand.Options.class + .getDeclaredField("rerankServiceOverride") + .getAnnotation(Schema.class) + .description(); + assertFalse(description.isBlank()); + var document = inputSchema("findAndRerank"); + assertEquals( + description, + schemaMetadata(document, propertySchema(document, "options", "rerank"), "description") + .asText()); + } + + /** Schema acceptance must agree with service validation of the override fields. */ + @ParameterizedTest + @MethodSource("rerankOverrides") + void rerankOverrideSchemaMatchesServiceValidation(String json, boolean accepted) + throws Exception { + var schema = inputSchema("findAndRerank"); + var payload = objectMapper.readTree(json); + var command = toolInput("{}"); + command.putObject("options").set("rerank", payload); + var errors = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012) + .getSchema(schema) + .validate(command); + assertEquals( + accepted, errors.isEmpty(), () -> "Schema validation for " + json + " returned " + errors); + + Executable validate = () -> validateRerankOverride(payload, AuthenticationType.NONE); + if (accepted) { + assertDoesNotThrow(validate); + } else { + var error = assertThrows(RequestException.class, validate); + assertEquals(RequestException.Code.INVALID_RERANK_OVERRIDE.name(), error.code); + } + } + + static Stream rerankOverrides() { + return Stream.of( + Arguments.of( + """ + {"provider": "nvidia", "modelName": "nvidia/rerank-v1"} + """, + true), + Arguments.of( + """ + {"provider": "nvidia", "modelName": "nvidia/rerank-v1", + "authentication": {}, "parameters": {}} + """, + true), + Arguments.of( + """ + {"provider": "nvidia", "modelName": "nvidia/rerank-v1", + "authentication": null, "parameters": null} + """, + true), + Arguments.of( + """ + {"provider": "nvidia", "modelName": "nvidia/rerank-v1", + "parameters": {"truncate": "END"}} + """, + false), + Arguments.of("{}", false), + Arguments.of( + """ + {"provider": "nvidia"} + """, + false), + Arguments.of( + """ + {"modelName": "nvidia/rerank-v1"} + """, + false), + Arguments.of( + """ + {"provider": null, "modelName": "nvidia/rerank-v1"} + """, + false), + Arguments.of( + """ + {"provider": "nvidia", "modelName": null} + """, + false), + Arguments.of( + """ + {"provider": null, "modelName": null} + """, + false)); + } + + /** The schema allows credentials because the provider configuration determines support. */ + @ParameterizedTest + @EnumSource(AuthenticationType.class) + void rerankAuthenticationValidationDependsOnProvider(AuthenticationType authenticationType) + throws Exception { + var payload = + objectMapper.readTree( + """ + {"provider": "nvidia", "modelName": "nvidia/rerank-v1", + "authentication": {"providerKey": "credential-reference"}} + """); + var command = toolInput("{}"); + command.putObject("options").set("rerank", payload); + assertSchemaAccepts(inputSchema("findAndRerank"), command); + + Executable validate = () -> validateRerankOverride(payload, authenticationType); + if (authenticationType == AuthenticationType.SHARED_SECRET) { + assertDoesNotThrow(validate); + } else { + var error = assertThrows(RequestException.class, validate); + assertEquals(RequestException.Code.INVALID_RERANK_OVERRIDE.name(), error.code); + } + } + + private CollectionRerankDef.RerankServiceDef validateRerankOverride( + JsonNode payload, AuthenticationType authenticationType) throws Exception { + var override = + objectMapper.treeToValue(payload, CreateCollectionCommand.Options.RerankServiceDesc.class); + return CollectionRerankDef.validateServiceDesc( + supportedRerankingProvider(authenticationType), + override.provider(), + override.modelName(), + override.authentication(), + override.parameters(), + RequestException.Code.INVALID_RERANK_OVERRIDE); + } + + private static RerankingProvidersConfig supportedRerankingProvider( + AuthenticationType authenticationType) { + var model = + new ModelConfigImpl( + "nvidia/rerank-v1", + new ApiModelSupport.ApiModelSupportImpl( + ApiModelSupport.SupportStatus.SUPPORTED, Optional.empty()), + false, + "https://example.com/rerank", + new RequestPropertiesImpl(3, 10, 100, 100, 0.5, 10)); + return new RerankingProvidersConfigImpl( + Map.of( + "nvidia", + new RerankingProviderConfigImpl( + false, + "nvidia", + true, + Map.of(authenticationType, new AuthenticationConfigImpl(true, List.of())), + List.of(model)))); + } + + @ParameterizedTest + @ValueSource(strings = {"50", "{\"$vector\": 50, \"$lexical\": 10}"}) + void rerankToolAcceptsWirePayloads(String limits) throws Exception { + var payload = + toolInput( + """ + { + "filter": {"age": {"$gt": 40}}, + "sort": {"$hybrid": "search text"}, + "options": {"rerank": {"provider": "nvidia", "modelName": "nvidia/rerank-v1"}} + } + """); + ((ObjectNode) payload.path("options")).set("hybridLimits", objectMapper.readTree(limits)); + assertSchemaAccepts(inputSchema("findAndRerank"), payload); + } + + @ParameterizedTest + @ValueSource( + strings = { + "{\"$hybrid\":{\"$vector\":null,\"$lexical\":\"query\"}}", + "{\"$hybrid\":{\"$vectorize\":null,\"$lexical\":\"query\"}}", + "{\"$hybrid\":{\"$vectorize\":\"query\",\"$lexical\":null}}", + "{\"$hybrid\":{\"$vector\":[1,2,3],\"$lexical\":null}}", + "{\"$hybrid\":{\"$vector\":null}}", + "{\"$hybrid\":{\"$vectorize\":null}}", + "{\"$hybrid\":{\"$lexical\":null}}", + "{\"$hybrid\":{}}", + "{}", + "null" + }) + void sortSchemaAcceptsNullableAndEmptyWireForms(String json) throws Exception { + var sort = objectMapper.readTree(json); + assertSchemaAccepts(schemaFor(FindAndRerankSort.class), sort); + var clause = assertDoesNotThrow(() -> objectMapper.treeToValue(sort, FindAndRerankSort.class)); + if (!sort.isNull()) { + assertNotNull(clause); + } + + var payload = toolInput("{}"); + payload.set("sort", sort); + assertSchemaAccepts(inputSchema("findAndRerank"), payload); + } + + @Test + void rerankOverrideExamplesDeserialize() throws Exception { + var document = inputSchema("findAndRerank"); + var rerank = propertySchema(document, "options", "rerank"); + var examples = schemaMetadata(document, rerank, "examples"); + assertFalse(examples.isMissingNode()); + assertFalse(examples.isEmpty()); + for (var example : examples) { + var command = toolInput("{}"); + command.putObject("options").set("rerank", example); + assertSchemaAccepts(document, command); + assertNotNull( + objectMapper.treeToValue( + example, CreateCollectionCommand.Options.RerankServiceDesc.class)); + } + } + + @Test + void findToolAcceptsWirePayloads() throws Exception { + var payload = + toolInput( + """ + {"filter": {"age": {"$gt": 40}}, "sort": {"user.age": -1, "user.name": 1}} + """); + assertSchemaAccepts(inputSchema("find"), payload); + } + + @ParameterizedTest + @MethodSource("invalidToolPayloads") + void toolSchemasRejectInvalidPayloads(String tool, String json) throws Exception { + var schema = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012).getSchema(inputSchema(tool)); + assertFalse(schema.validate(toolInput(json)).isEmpty()); + } + + static Stream invalidToolPayloads() { + return Stream.of( + Arguments.of("find", "{\"filter\": 5}"), + Arguments.of("find", "{\"sort\": []}"), + Arguments.of("findAndRerank", "{\"sort\": {\"vectorizeSort\": \"query\"}}"), + Arguments.of("findAndRerank", "{\"sort\": {\"$hybrid\": {\"$vector\": \"query\"}}}"), + Arguments.of("findAndRerank", "{\"options\": {\"hybridLimits\": {\"$vector\": 50}}}"), + Arguments.of("findAndRerank", "{\"options\": {\"rerank\": {\"provider\": \"nvidia\"}}}")); + } + + private ObjectNode toolInput(String json) throws Exception { + // Routing arguments are required by MCP but are not part of the command clauses. + return ((ObjectNode) objectMapper.readTree(json)) + .put("keyspace", "default_keyspace") + .put("collection", "documents"); + } + + /** Positive examples alone would also pass against a schema that accepts every JSON value. */ + @ParameterizedTest + @MethodSource("invalidWireShapes") + void schemaRejectsInvalidWireShapes(Class clauseType, String json) throws Exception { + var schema = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012) + .getSchema(schemaFor(clauseType)); + assertFalse(schema.validate(objectMapper.readTree(json)).isEmpty()); + } + + static Stream invalidWireShapes() { + return Stream.of( + Arguments.of(FindAndRerankSort.class, "{\"vectorizeSort\": \"query\"}"), + Arguments.of(FindAndRerankSort.class, "{\"$hybrid\": {\"$lexical\": 1}}"), + Arguments.of(FindAndRerankSort.class, "{\"$hybrid\": {\"$vector\": \"query\"}}"), + Arguments.of( + FindAndRerankSort.class, "{\"$hybrid\": {\"$vector\": null, \"$vectorize\": null}}"), + Arguments.of( + FindAndRerankCommand.HybridLimits.class, "{\"vectorLimit\": 50, \"lexicalLimit\": 10}"), + Arguments.of(FindAndRerankCommand.HybridLimits.class, "{\"$vector\": 50}"), + Arguments.of( + FindAndRerankCommand.HybridLimits.class, "{\"$vector\": 50, \"$lexical\": \"10\"}")); + } + + private void assertSchemaAccepts(JsonNode schema, JsonNode example) { + var validator = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012).getSchema(schema); + assertTrue(validator.validate(example).isEmpty(), () -> validator.validate(example).toString()); + } + + @Test + void collectionCreationAcceptsRerankAuthenticationAndParameters() throws Exception { + var payload = + objectMapper.readTree( + """ + { + "provider": "nvidia", "modelName": "nvidia/rerank-v1", + "authentication": {"providerKey": "credential-reference"}, + "parameters": {"truncate": "END"} + } + """); + assertSchemaAccepts( + schemaFor(CreateCollectionCommand.Options.RerankServiceDesc.class), payload); + assertEquals( + new CreateCollectionCommand.Options.RerankServiceDesc( + "nvidia", + "nvidia/rerank-v1", + Map.of("providerKey", "credential-reference"), + Map.of("truncate", "END")), + objectMapper.treeToValue(payload, CreateCollectionCommand.Options.RerankServiceDesc.class)); + } + + @Test + void allRegisteredToolsGenerateInputSchemas() { + var count = 0; + // Exercise named OpenAPI lookups for every tool, including ordinary options records. + for (var tool : toolManager) { + var schema = + assertDoesNotThrow( + () -> objectMapper.readTree(schemaGenerator.generate(tool).asJson()), tool.name()); + assertTrue(schema.path("properties").isObject(), tool.name()); + count++; + } + assertTrue(count > 0, "The MCP server must register tools"); + } + + private JsonNode inputSchema(String name) throws Exception { + // The registered tool includes nested command fields and the application's customizer wiring. + var tool = toolManager.getTool(name); + assertNotNull(tool); + return objectMapper.readTree(schemaGenerator.generate(tool).asJson()); + } + + private JsonNode schemaFor(Class type) { + var builder = + new SchemaGeneratorConfigBuilder( + objectMapper, SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON) + .without(Option.SCHEMA_VERSION_INDICATOR); + new McpSchemaCustomizer().customize(builder); + var schema = new SchemaGenerator(builder.build()).generateSchema(type); + McpOpenApiSchemaAdapter.bundleOpenApiSchemas( + schema, McpOpenApiSchemaAdapter.readGeneratedOpenApiDocument()); + return schema; + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpOpenApiSchemaAdapterTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpOpenApiSchemaAdapterTest.java new file mode 100644 index 0000000000..b093b66e99 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpOpenApiSchemaAdapterTest.java @@ -0,0 +1,352 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.networknt.schema.JsonSchemaFactory; +import com.networknt.schema.SpecVersion; +import org.junit.jupiter.api.Test; + +/** Checks that standalone MCP schemas retain OpenAPI references and validation semantics. */ +class McpOpenApiSchemaAdapterTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + void shouldCopyOnlyReachableComponentsAndRebaseTheirReferences() throws Exception { + var schema = bundle("{\"$ref\":\"#/components/schemas/ObjectValue\"}"); + + assertEquals("#/$defs/_openapi_ObjectValue", schema.path("$ref").asText()); + var definitions = schema.path("$defs"); + assertEquals(2, definitions.size()); + assertTrue(definitions.has("_openapi_ObjectValue")); + assertTrue(definitions.has("_openapi_Value")); + assertFalse(definitions.has("_openapi_Unused")); + assertEquals( + "#/$defs/_openapi_Value", + definitions + .path("_openapi_ObjectValue") + .path("properties") + .path("a") + .path("$ref") + .asText()); + assertAccepts(schema, "{\"a\":10}"); + assertRejects(schema, "{\"a\":15}"); + } + + @Test + void shouldLeaveFallbackSchemasAndExistingDefinitionsUnchanged() throws Exception { + var input = + MAPPER.readTree( + """ + { + "type":"object", + "properties":{"value":{"$ref":"#/$defs/Local"}}, + "$defs":{"Local":{"type":"string"}} + } + """); + + var adapted = bundle(input); + + assertEquals(input, adapted); + assertAccepts(adapted, "{\"value\":\"text\"}"); + assertRejects(adapted, "{\"value\":5}"); + } + + @Test + void shouldResolveFallbackDefinitionsWithinEachToolArgument() throws Exception { + var input = + MAPPER.readTree( + """ + { + "type":"object", + "properties":{ + "options/~":{ + "type":"object", + "properties":{ + "value":{"$ref":"#/$defs/Local"}, + "shared":{"$ref":"#/components/schemas/Value"} + }, + "$defs":{"Local":{ + "type":"object", + "properties":{ + "name":{"type":"string"}, + "next":{"$ref":"#/$defs/Local"} + } + }}, + "examples":[{"value":{"$ref":"#/$defs/Local"}}] + }, + "other":{ + "type":"object", + "properties":{"value":{"$ref":"#/$defs/Local"}}, + "$defs":{"Local":{"type":"integer"}} + } + } + } + """); + var schema = bundle(input); + + assertEquals( + "#/properties/options~1~0/$defs/Local", + schema.at("/properties/options~1~0/properties/value/$ref").asText()); + assertEquals( + input.at("/properties/options~1~0/examples"), + schema.at("/properties/options~1~0/examples")); + assertAccepts( + schema, + """ + {"options/~":{"value":{"name":"first","next":{"name":"second"}},"shared":10}, + "other":{"value":2}} + """); + assertRejects(schema, "{\"options/~\":{\"value\":{\"next\":{\"name\":2}}}}"); + assertRejects(schema, "{\"options/~\":{\"shared\":15}}"); + assertRejects(schema, "{\"other\":{\"value\":\"wrong\"}}"); + } + + @Test + void shouldPreserveValidRootReferencesWhenArgumentDefinitionsHaveTheSameName() throws Exception { + var input = + MAPPER.readTree( + """ + { + "type":"object", + "properties":{"options":{ + "type":"object", + "properties":{"value":{"$ref":"#/$defs/Local"}}, + "$defs":{"Local":{"type":"integer"}} + }}, + "$defs":{"Local":{"type":"string"}} + } + """); + var schema = bundle(input); + + assertEquals(input, schema); + assertAccepts(schema, "{\"options\":{\"value\":\"text\"}}"); + assertRejects(schema, "{\"options\":{\"value\":5}}"); + } + + @Test + void shouldRetainExistingDefinitionsWhenAddingOpenApiComponents() throws Exception { + var schema = + bundle( + """ + { + "type":"object", + "properties":{ + "local":{"$ref":"#/$defs/Local"}, + "remote":{"$ref":"#/components/schemas/Value"} + }, + "$defs":{"Local":{"type":"string"}} + } + """); + + assertEquals("#/$defs/Local", schema.path("properties").path("local").path("$ref").asText()); + assertEquals("string", schema.path("$defs").path("Local").path("type").asText()); + assertTrue(schema.path("$defs").has("_openapi_Value")); + assertAccepts(schema, "{\"local\":\"text\",\"remote\":10}"); + assertRejects(schema, "{\"local\":\"text\",\"remote\":15}"); + } + + @Test + void shouldSupportRecursiveComponentsWithoutExpandingThem() throws Exception { + var schema = bundle("{\"$ref\":\"#/components/schemas/Recursive\"}"); + + assertEquals(1, schema.path("$defs").size()); + assertEquals( + "#/$defs/_openapi_Recursive", + schema + .path("$defs") + .path("_openapi_Recursive") + .path("properties") + .path("next") + .path("anyOf") + .get(0) + .path("$ref") + .asText()); + assertAccepts(schema, "{\"value\":1,\"next\":{\"value\":2,\"next\":null}}"); + assertRejects(schema, "{\"value\":1,\"next\":{\"value\":\"wrong\"}}"); + } + + @Test + void shouldPreserveEscapedComponentNamesAndNestedPointerSuffixes() throws Exception { + var document = fixtureDocument(); + ((ObjectNode) document.path("components").path("schemas")) + .putObject("A/B~C") + .put("type", "object") + .putObject("properties") + .putObject("value") + .put("type", "integer"); + var input = + MAPPER.readTree( + """ + { + "type":"object", + "properties":{ + "whole":{"$ref":"#/components/schemas/A~1B~0C"}, + "value":{"$ref":"#/components/schemas/A~1B~0C/properties/value"} + } + } + """); + + var schema = bundle(input, document); + + assertTrue(schema.path("$defs").has("_openapi_A/B~C")); + assertEquals( + "#/$defs/_openapi_A~1B~0C", schema.path("properties").path("whole").path("$ref").asText()); + assertEquals( + "#/$defs/_openapi_A~1B~0C/properties/value", + schema.path("properties").path("value").path("$ref").asText()); + assertAccepts(schema, "{\"whole\":{\"value\":1},\"value\":2}"); + assertRejects(schema, "{\"whole\":{\"value\":1},\"value\":\"wrong\"}"); + } + + @Test + void shouldRetainOverlappingConstraintsOnReferenceSiblings() throws Exception { + var schema = bundle("{\"$ref\":\"#/components/schemas/Value\",\"maximum\":20}"); + + assertEquals(20, schema.path("maximum").asInt()); + assertAccepts(schema, "10"); + assertRejects(schema, "15"); + } + + @Test + void shouldRetainDisjointConstraintsOnReferenceSiblings() throws Exception { + var schema = + bundle("{\"$ref\":\"#/components/schemas/ObjectValue\",\"additionalProperties\":false}"); + + assertTrue(schema.has("$ref")); + assertFalse(schema.path("additionalProperties").asBoolean()); + assertAccepts(schema, "{}"); + assertRejects(schema, "{\"a\":1}"); + } + + @Test + void shouldKeepReferenceSiblingsInsideCopiedComponents() throws Exception { + var document = fixtureDocument(); + ((ObjectNode) + document + .path("components") + .path("schemas") + .path("ObjectValue") + .path("properties") + .path("a")) + .put("maximum", 20); + var schema = + bundle(MAPPER.readTree("{\"$ref\":\"#/components/schemas/ObjectValue\"}"), document); + + assertAccepts(schema, "{\"a\":10}"); + assertRejects(schema, "{\"a\":15}"); + } + + @Test + void shouldPreserveAPropertyNamedRefAndLiteralReferencesInInstanceData() throws Exception { + var input = + MAPPER.readTree( + """ + { + "type":"object", + "properties":{"$ref":{"$ref":"#/components/schemas/Value"}}, + "examples":[{"$ref":"#/components/schemas/Unknown"}], + "default":{"$ref":"#/components/schemas/Unknown"}, + "const":{"$ref":"#/components/schemas/Unknown"}, + "enum":[{"$ref":"#/components/schemas/Unknown"}] + } + """); + var schema = bundle(input); + + assertEquals( + "#/$defs/_openapi_Value", schema.path("properties").path("$ref").path("$ref").asText()); + for (var keyword : new String[] {"examples", "default", "const", "enum"}) { + assertEquals(input.path(keyword), schema.path(keyword)); + } + } + + @Test + void shouldChangeTheGivenSchemaInPlace() throws Exception { + var schema = (ObjectNode) MAPPER.readTree("{\"$ref\":\"#/components/schemas/Value\"}"); + + McpOpenApiSchemaAdapter.bundleOpenApiSchemas(schema, fixtureDocument()); + + assertEquals("#/$defs/_openapi_Value", schema.path("$ref").asText()); + assertTrue(schema.path("$defs").has("_openapi_Value")); + } + + @Test + void shouldNotChangeTheOpenApiDocumentOrShareComponentCopies() throws Exception { + var document = fixtureDocument(); + var originalDocument = document.deepCopy(); + var input = MAPPER.readTree("{\"$ref\":\"#/components/schemas/Value\"}"); + var schema = bundle(input, document); + + ((ObjectNode) schema.path("$defs").path("_openapi_Value")).put("maximum", 100); + + assertEquals(originalDocument, document); + assertRejects(bundle(input, document), "15"); + } + + @Test + void shouldRejectMissingComponentReferences() throws Exception { + var input = MAPPER.readTree("{\"$ref\":\"#/components/schemas/Missing\"}"); + + assertThrows(IllegalStateException.class, () -> bundle(input)); + } + + @Test + void shouldEscapeComponentNamesInReferences() { + assertEquals("#/components/schemas/A~1B~0C", McpOpenApiSchemaAdapter.componentRef("A/B~C")); + } + + /** Bundles a copy of {@code input}, so tests can compare the result with the input. */ + private static ObjectNode bundle(JsonNode input, JsonNode document) { + ObjectNode schema = input.deepCopy(); + McpOpenApiSchemaAdapter.bundleOpenApiSchemas(schema, document); + return schema; + } + + private static ObjectNode bundle(JsonNode input) throws Exception { + return bundle(input, fixtureDocument()); + } + + private static ObjectNode bundle(String input) throws Exception { + return bundle(MAPPER.readTree(input)); + } + + private static void assertAccepts(JsonNode schema, String input) throws Exception { + var errors = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012) + .getSchema(schema) + .validate(MAPPER.readTree(input)); + assertTrue(errors.isEmpty(), errors::toString); + } + + private static void assertRejects(JsonNode schema, String input) throws Exception { + var errors = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012) + .getSchema(schema) + .validate(MAPPER.readTree(input)); + assertFalse(errors.isEmpty()); + } + + private static JsonNode fixtureDocument() throws Exception { + return MAPPER.readTree( + """ + { + "components":{"schemas":{ + "Value":{"type":"integer","maximum":10}, + "ObjectValue":{"type":"object","properties":{"a":{"$ref":"#/components/schemas/Value"}}}, + "Recursive":{ + "type":"object", + "properties":{ + "value":{"type":"integer"}, + "next":{"anyOf":[{"$ref":"#/components/schemas/Recursive"},{"type":"null"}]} + }, + "required":["value"], + "additionalProperties":false + }, + "Unused":{"type":"string"} + }} + } + """); + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaCustomizerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaCustomizerTest.java new file mode 100644 index 0000000000..9ec348f1f1 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaCustomizerTest.java @@ -0,0 +1,194 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.github.victools.jsonschema.generator.Option; +import com.github.victools.jsonschema.generator.OptionPreset; +import com.github.victools.jsonschema.generator.SchemaGenerator; +import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder; +import com.github.victools.jsonschema.generator.SchemaVersion; +import org.eclipse.microprofile.openapi.annotations.media.Schema; +import org.junit.jupiter.api.Test; + +/** + * Checks OpenAPI references and annotation descriptions supplied by {@link McpSchemaCustomizer}. + */ +class McpSchemaCustomizerTest { + + /** + * Test record that mimics nested command structures such as {@link + * io.stargate.sgv2.jsonapi.api.model.command.impl.CreateTableCommand}. CreateTableCommand → + * TableDefinitionDesc → PrimaryKeyDesc + */ + record TestPrimaryKey( + @Schema(description = "Partition key columns") String[] keys, + @Schema(description = "Clustering key ordering") String ordering) {} + + record TestDefinition( + @Schema(description = "Column definitions for the table") String columns, + @Schema(description = "Primary key definition") TestPrimaryKey primaryKey) {} + + record TestOptions( + @Schema(description = "Flag to ignore if already exists") Boolean ifNotExists) {} + + record TestCreateTable( + @Schema(description = "Table definition") TestDefinition definition, + @Schema(description = "Configuration options") TestOptions options, + Boolean noAnnotation) {} + + @Schema(description = "Reusable value description") + record TestDescribedType(String value) {} + + record TestTypeDescription(TestDescribedType value) {} + + record TestHiddenFields(@Schema(hidden = true) String internal, String visible) {} + + record TestNamedField(@Schema(name = "aPropertyName") String value) {} + + @Schema(name = "CustomizerTestSchema") + record TestSharedType(String internalValue) {} + + @Schema(ref = "#/components/schemas/CustomizerTestSchema") + record TestReferencedType(String internalValue) {} + + record TestFieldValue(String internalValue) {} + + record TestSharedFields( + @Schema(ref = "#/components/schemas/CustomizerTestSchema", description = "First field") + TestFieldValue first, + @Schema(ref = "#/components/schemas/CustomizerTestSchema", description = "Second field") + TestFieldValue second) {} + + @Test + void shouldReferenceNamedComponentInsteadOfInternalJavaFields() { + var schema = schemaFor(TestSharedType.class); + + assertEquals("#/components/schemas/CustomizerTestSchema", schema.path("$ref").asText()); + assertFalse(schema.path("properties").has("internalValue")); + } + + @Test + void shouldKeepExplicitComponentReference() { + assertEquals(schemaFor(TestSharedType.class), schemaFor(TestReferencedType.class)); + } + + @Test + void shouldKeepFieldReferencesAndDescriptions() { + var properties = schemaFor(TestSharedFields.class).path("properties"); + + for (var name : new String[] {"first", "second"}) { + assertEquals( + "#/components/schemas/CustomizerTestSchema", properties.path(name).path("$ref").asText()); + assertFalse(properties.path(name).path("properties").has("internalValue")); + } + assertEquals("First field", properties.path("first").path("description").asText()); + assertEquals("Second field", properties.path("second").path("description").asText()); + assertFalse(schemaFor(TestSharedType.class).has("description")); + } + + @Test + void shouldOmitFieldsHiddenBySchemaAnnotation() { + var properties = schemaFor(TestHiddenFields.class).path("properties"); + + assertFalse(properties.has("internal")); + assertTrue(properties.has("visible")); + } + + @Test + void shouldNotTreatFieldSchemaNameAsComponentName() { + var schema = assertDoesNotThrow(() -> schemaFor(TestNamedField.class)); + + assertEquals("string", schema.path("properties").path("value").path("type").asText()); + } + + private static JsonNode schemaFor(Class type) { + var builder = baseConfigBuilder(); + new McpSchemaCustomizer().customize(builder); + return new SchemaGenerator(builder.build()).generateSchema(type); + } + + @Test + void shouldIncludeRootTypeDescription() { + var configBuilder = baseConfigBuilder(); + new McpSchemaCustomizer().customize(configBuilder); + + var schema = new SchemaGenerator(configBuilder.build()).generateSchema(TestDescribedType.class); + + assertEquals("Reusable value description", schema.path("description").asText()); + } + + @Test + void shouldIncludeNestedTypeDescriptionWithoutFieldDescription() { + var configBuilder = baseConfigBuilder(); + new McpSchemaCustomizer().customize(configBuilder); + + var schema = + new SchemaGenerator(configBuilder.build()).generateSchema(TestTypeDescription.class); + + assertEquals( + "Reusable value description", + schema.path("properties").path("value").path("description").asText()); + } + + @Test + void shouldIncludeSchemaDescriptionWithCustomizer() { + var configBuilder = baseConfigBuilder(); + new McpSchemaCustomizer().customize(configBuilder); + + var schema = new SchemaGenerator(configBuilder.build()).generateSchema(TestCreateTable.class); + var props = schema.path("properties"); + + // Level 1: top-level fields + assertEquals("Table definition", props.path("definition").path("description").asText()); + assertEquals("Configuration options", props.path("options").path("description").asText()); + assertTrue( + props.path("noAnnotation").path("description").isMissingNode(), + "Fields without @Schema should have no description"); + + // Level 2: nested inside TestDefinition + var defProps = props.path("definition").path("properties"); + assertEquals( + "Column definitions for the table", defProps.path("columns").path("description").asText()); + assertEquals( + "Primary key definition", defProps.path("primaryKey").path("description").asText()); + + // Level 2: nested inside TestOptions + var optProps = props.path("options").path("properties"); + assertEquals( + "Flag to ignore if already exists", + optProps.path("ifNotExists").path("description").asText()); + + // Level 3: nested inside TestPrimaryKey + var pkProps = defProps.path("primaryKey").path("properties"); + assertEquals("Partition key columns", pkProps.path("keys").path("description").asText()); + assertEquals("Clustering key ordering", pkProps.path("ordering").path("description").asText()); + } + + @Test + void shouldNotIncludeDescriptionWithoutCustomizer() { + var schema = + new SchemaGenerator(baseConfigBuilder().build()).generateSchema(TestCreateTable.class); + + // Spot-check: a deeply nested field should have no description + var pkProps = + schema + .get("properties") + .get("definition") + .get("properties") + .get("primaryKey") + .get("properties"); + assertNull( + pkProps.get("keys").get("description"), + "Without customizer, PLAIN_JSON preset should not resolve @Schema descriptions"); + } + + /** + * Creates a {@link SchemaGeneratorConfigBuilder} matching the configuration used by Quarkus MCP + * Server's {@code DefaultSchemaGenerator}. + */ + private static SchemaGeneratorConfigBuilder baseConfigBuilder() { + return new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON) + .without(Option.SCHEMA_VERSION_INDICATOR); + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaDescriptionCustomizerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaDescriptionCustomizerTest.java deleted file mode 100644 index 8ec17c71c1..0000000000 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaDescriptionCustomizerTest.java +++ /dev/null @@ -1,99 +0,0 @@ -package io.stargate.sgv2.jsonapi.api.v1.mcp; - -import static org.junit.jupiter.api.Assertions.*; - -import com.github.victools.jsonschema.generator.Option; -import com.github.victools.jsonschema.generator.OptionPreset; -import com.github.victools.jsonschema.generator.SchemaGenerator; -import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder; -import com.github.victools.jsonschema.generator.SchemaVersion; -import org.eclipse.microprofile.openapi.annotations.media.Schema; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link McpSchemaDescriptionCustomizer}. Verifies that the customizer correctly - * resolves {@link Schema#description()} into the generated JSON schema. - */ -class McpSchemaDescriptionCustomizerTest { - - /** - * Test record that mimics nested command structures such as {@link - * io.stargate.sgv2.jsonapi.api.model.command.impl.CreateTableCommand}. CreateTableCommand → - * TableDefinitionDesc → PrimaryKeyDesc - */ - record TestPrimaryKey( - @Schema(description = "Partition key columns") String[] keys, - @Schema(description = "Clustering key ordering") String ordering) {} - - record TestDefinition( - @Schema(description = "Column definitions for the table") String columns, - @Schema(description = "Primary key definition") TestPrimaryKey primaryKey) {} - - record TestOptions( - @Schema(description = "Flag to ignore if already exists") Boolean ifNotExists) {} - - record TestCreateTable( - @Schema(description = "Table definition") TestDefinition definition, - @Schema(description = "Configuration options") TestOptions options, - Boolean noAnnotation) {} - - @Test - void shouldIncludeSchemaDescriptionWithCustomizer() { - var configBuilder = baseConfigBuilder(); - new McpSchemaDescriptionCustomizer().customize(configBuilder); - - var schema = new SchemaGenerator(configBuilder.build()).generateSchema(TestCreateTable.class); - var props = schema.get("properties"); - - // Level 1: top-level fields - assertEquals("Table definition", props.get("definition").get("description").asText()); - assertEquals("Configuration options", props.get("options").get("description").asText()); - assertNull( - props.get("noAnnotation").get("description"), - "Fields without @Schema should have no description"); - - // Level 2: nested inside TestDefinition - var defProps = props.get("definition").get("properties"); - assertEquals( - "Column definitions for the table", defProps.get("columns").get("description").asText()); - assertEquals("Primary key definition", defProps.get("primaryKey").get("description").asText()); - - // Level 2: nested inside TestOptions - var optProps = props.get("options").get("properties"); - assertEquals( - "Flag to ignore if already exists", - optProps.get("ifNotExists").get("description").asText()); - - // Level 3: nested inside TestPrimaryKey - var pkProps = defProps.get("primaryKey").get("properties"); - assertEquals("Partition key columns", pkProps.get("keys").get("description").asText()); - assertEquals("Clustering key ordering", pkProps.get("ordering").get("description").asText()); - } - - @Test - void shouldNotIncludeDescriptionWithoutCustomizer() { - var schema = - new SchemaGenerator(baseConfigBuilder().build()).generateSchema(TestCreateTable.class); - - // Spot-check: a deeply nested field should have no description - var pkProps = - schema - .get("properties") - .get("definition") - .get("properties") - .get("primaryKey") - .get("properties"); - assertNull( - pkProps.get("keys").get("description"), - "Without customizer, PLAIN_JSON preset should not resolve @Schema descriptions"); - } - - /** - * Creates a {@link SchemaGeneratorConfigBuilder} matching the configuration used by Quarkus MCP - * Server's {@code DefaultSchemaGenerator}. - */ - private static SchemaGeneratorConfigBuilder baseConfigBuilder() { - return new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON) - .without(Option.SCHEMA_VERSION_INDICATOR); - } -} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaIntegrationTest.java new file mode 100644 index 0000000000..6ad33411b8 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSchemaIntegrationTest.java @@ -0,0 +1,170 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.networknt.schema.JsonSchema; +import com.networknt.schema.JsonSchemaFactory; +import com.networknt.schema.SpecVersion; +import io.quarkiverse.mcp.server.test.McpAssured.ToolsPage; +import io.quarkus.test.junit.QuarkusIntegrationTest; +import io.quarkus.test.junit.TestProfile; +import io.stargate.sgv2.jsonapi.testresource.NoGlobalResourcesTestProfile; +import io.vertx.core.json.JsonObject; +import org.junit.jupiter.api.Test; + +/** Checks the schemas clients receive from the packaged application's tools/list endpoint. */ +@QuarkusIntegrationTest +@TestProfile(NoGlobalResourcesTestProfile.Impl.class) +class McpSchemaIntegrationTest extends McpIntegrationTestBase { + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final JsonSchemaFactory SCHEMAS = + JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012); + + // Tool discovery loads the packaged OpenAPI document without executing database commands. + @Test + void advertisedToolSchemasResolveTheirReferences() { + mcpClient + .when() + .toolsList( + page -> { + assertFalse(page.tools().isEmpty()); + assertNull(page.nextCursor(), "Validate every advertised tool in this response"); + assertNotNull(page.findByName("find")); + assertNotNull(page.findByName("findAndRerank")); + for (var tool : page.tools()) { + assertEquals("object", tool.inputSchema().getString("type"), tool.name()); + assertDoesNotThrow( + () -> schemaFor(page, tool.name()).initializeValidators(), tool.name()); + } + }) + .thenAssertResults(); + } + + @Test + void findAdvertisesDocumentPathsInsteadOfJavaFields() { + mcpClient + .when() + .toolsList( + page -> { + var schema = schemaFor(page, "find"); + var arguments = + new JsonObject( + """ + { + "keyspace":"default_keyspace", + "collection":"schema_test", + "filter":{"empty":{"$eq":true},"json":"example"}, + "sort":{"age":-1} + } + """); + assertAccepts(schema, arguments); + assertRejects(schema, arguments.copy().put("filter", "not an object")); + assertRejects(schema, arguments.copy().put("sort", "not an object")); + }) + .thenAssertResults(); + } + + @Test + void createCollectionResolvesProviderOptionsDefinitions() { + mcpClient + .when() + .toolsList( + page -> { + var schema = schemaFor(page, "createCollection"); + var arguments = + new JsonObject( + """ + { + "keyspace":"default_keyspace", + "collection":"schema_test", + "options":{"rerank":{"service":{ + "provider":"nvidia", + "modelName":"nvidia/rerank-v1", + "authentication":{"providerKey":"credential-reference"}, + "parameters":{"truncate":"END"} + }}} + } + """); + assertAccepts(schema, arguments); + + var invalidAuthentication = arguments.copy(); + invalidAuthentication + .getJsonObject("options") + .getJsonObject("rerank") + .getJsonObject("service") + .put("authentication", "not an object"); + assertRejects(schema, invalidAuthentication); + }) + .thenAssertResults(); + } + + @Test + void findAndRerankPreservesNestedConstraints() { + mcpClient + .when() + .toolsList( + page -> { + var schema = schemaFor(page, "findAndRerank"); + var arguments = + new JsonObject( + """ + { + "keyspace":"default_keyspace", + "collection":"schema_test", + "filter":null, + "sort":{"$hybrid":{"$vector":[1,2,3],"$lexical":"coffee"}}, + "options":{"hybridLimits":{"$vector":50,"$lexical":10}} + } + """); + assertAccepts(schema, arguments); + + var binaryVector = arguments.copy(); + binaryVector + .getJsonObject("sort") + .getJsonObject("$hybrid") + .put("$vector", new JsonObject().put("$binary", "P4AAAEAAAABAQAAA")); + assertAccepts(schema, binaryVector); + + var shorthand = + arguments.copy().put("sort", new JsonObject().put("$hybrid", "coffee")); + shorthand.getJsonObject("options").put("hybridLimits", 50); + assertAccepts(schema, shorthand); + + var conflictingVector = arguments.copy(); + conflictingVector + .getJsonObject("sort") + .getJsonObject("$hybrid") + .put("$vectorize", "coffee"); + assertRejects(schema, conflictingVector); + + var missingLimit = arguments.copy(); + missingLimit + .getJsonObject("options") + .getJsonObject("hybridLimits") + .remove("$lexical"); + assertRejects(schema, missingLimit); + }) + .thenAssertResults(); + } + + private static JsonSchema schemaFor(ToolsPage page, String name) { + var tool = page.findByName(name); + assertNotNull(tool, name); + return SCHEMAS.getSchema(json(tool.inputSchema())); + } + + private static void assertAccepts(JsonSchema schema, JsonObject arguments) { + var errors = schema.validate(json(arguments)); + assertTrue(errors.isEmpty(), errors::toString); + } + + private static void assertRejects(JsonSchema schema, JsonObject arguments) { + assertFalse(schema.validate(json(arguments)).isEmpty()); + } + + private static JsonNode json(JsonObject value) { + return assertDoesNotThrow(() -> MAPPER.readTree(value.encode())); + } +}