diff --git a/mcp/mcp-schemas/model/main.smithy b/mcp/mcp-schemas/model/main.smithy index 0f1e5ac95..68b289722 100644 --- a/mcp/mcp-schemas/model/main.smithy +++ b/mcp/mcp-schemas/model/main.smithy @@ -126,6 +126,10 @@ structure JsonObjectSchema { description: String + /// Present when the object is a discriminated polymorphic type (see the smithy.mcp#oneOf + /// trait): the instance must additionally match exactly one of these variant schemas. + oneOf: JsonSchemaList + @jsonName("$schema") schema: String = "http://json-schema.org/draft-07/schema#" } diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java index 64ee01def..db3ec8103 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java @@ -71,16 +71,8 @@ McpToolDescriptor createTool(String serverId, Service service, Operation operati var info = ToolInfo.builder() .name(operationName) .description(createDescription(service.schema().id().getName(), operationName, operationSchema)) - .inputSchema(createObjectSchema( - operation.getApiOperation().inputSchema(), - operation.getApiOperation().inputSchema(), - new HashSet<>(), - cache)) - .outputSchema(createObjectSchema( - operation.getApiOperation().outputSchema(), - operation.getApiOperation().outputSchema(), - new HashSet<>(), - cache)) + .inputSchema(createRootSchema(operation.getApiOperation().inputSchema(), cache)) + .outputSchema(createRootSchema(operation.getApiOperation().outputSchema(), cache)) .annotations(createAnnotations(operationSchema)) .build(); return new McpToolDescriptor( @@ -138,7 +130,7 @@ private JsonObjectSchema createObjectSchema( if (child.hasTrait(TraitKey.REQUIRED_TRAIT)) { required.add(child.memberName()); } - properties.put(child.memberName(), Document.of(createMemberSchema(child, visited, cache))); + properties.put(child.memberName(), Document.of(createSchema(child, visited, cache))); } visited.remove(targetId); @@ -150,13 +142,46 @@ private JsonObjectSchema createObjectSchema( return (JsonObjectSchema) withDescription(result, memberDescription(member)); } + private JsonObjectSchema createRootSchema(Schema root, Map cache) { + return asJsonObjectSchema(root, createSchema(root, new HashSet<>(), cache)); + } + + /** + * Coerces a rendered root schema into the object-typed schema {@link ToolInfo} requires. A + * polymorphic root ({@link JsonOneOfSchema}) becomes an object constrained by the same + * {@code oneOf} variants; on the wire this only adds the {@code $schema} annotation, since + * {@link JsonOneOfSchema} already declares {@code "type": "object"}. An untyped document root + * becomes a permissive object. Anything else has no object representation and degrades to a + * permissive object schema with a warning rather than failing the entire tool listing. + */ + private static JsonObjectSchema asJsonObjectSchema(Schema root, SerializableShape schema) { + return switch (schema) { + case JsonObjectSchema object -> object; + case JsonOneOfSchema oneOf -> { + var builder = JsonObjectSchema.builder().oneOf(oneOf.getOneOf()); + if (oneOf.getDescription() != null) { + builder.description(oneOf.getDescription()); + } + yield builder.build(); + } + case JsonDocumentSchema document -> { + var builder = JsonObjectSchema.builder(); + if (document.getDescription() != null) { + builder.description(document.getDescription()); + } + yield builder.build(); + } + default -> JsonObjectSchema.builder().build(); + }; + } + private JsonArraySchema createArraySchema( Schema member, Schema target, Set visited, Map cache ) { - var items = createMemberSchema(target.listMember(), visited, cache); + var items = createSchema(target.listMember(), visited, cache); var itemDocument = target.hasTrait(TraitKey.SPARSE_TRAIT) ? Document.of(Map.of( "anyOf", @@ -233,10 +258,9 @@ private SerializableShape createOneOfSchema( var variants = new ArrayList(); for (var definition : oneOf.getMembers()) { - var target = schemaIndex.getSchema(definition.getTarget()); variants.add(createUnionVariant( definition.getName(), - createObjectSchema(target, target, visited, cache))); + createSchema(schemaIndex.getSchema(definition.getTarget()), visited, cache))); } visited.remove(targetId); @@ -264,7 +288,7 @@ private SerializableShape createUnionSchema( for (var child : target.members()) { variants.add(createUnionVariant( child.memberName(), - createMemberSchema(child, visited, cache))); + createSchema(child, visited, cache))); } visited.remove(targetId); @@ -273,18 +297,22 @@ private SerializableShape createUnionSchema( return withDescription(result, memberDescription(member)); } - private SerializableShape createMemberSchema( - Schema member, + /** + * Renders any schema, member or not, by dispatching on the type of the shape it resolves to. + */ + private SerializableShape createSchema( + Schema schema, Set visited, Map cache ) { - return switch (member.type()) { - case LIST, SET -> createArraySchema(member, member.memberTarget(), visited, cache); - case MAP -> createMapSchema(member, member.memberTarget(), visited, cache); - case STRUCTURE -> createObjectSchema(member, member.memberTarget(), visited, cache); - case UNION -> createUnionSchema(member, member.memberTarget(), visited, cache); - case DOCUMENT -> createDocumentSchema(member, visited, cache); - default -> createPrimitiveSchema(member); + var target = schema.isMember() ? schema.memberTarget() : schema; + return switch (target.type()) { + case LIST, SET -> createArraySchema(schema, target, visited, cache); + case MAP -> createMapSchema(schema, target, visited, cache); + case STRUCTURE -> createObjectSchema(schema, target, visited, cache); + case UNION -> createUnionSchema(schema, target, visited, cache); + case DOCUMENT -> createDocumentSchema(schema, visited, cache); + default -> createPrimitiveSchema(schema); }; } @@ -294,7 +322,7 @@ private JsonObjectSchema createMapSchema( Set visited, Map cache ) { - var value = createMemberSchema(target.mapValueMember(), visited, cache); + var value = createSchema(target.mapValueMember(), visited, cache); var additionalProperties = target.hasTrait(TraitKey.SPARSE_TRAIT) ? Document.of(Map.of( "anyOf", diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpServerTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpServerTest.java index 4a7dfd48e..acc6ef4bb 100644 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpServerTest.java +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpServerTest.java @@ -1810,6 +1810,110 @@ private void writeNotification(String method, Document params) { .assemble() .unwrap(); + private static final String ONE_OF_ROOT_MODEL_STR = + """ + $version: "2" + + namespace smithy.test.oneofroot + + use smithy.mcp#oneOf + + // One service per operation so each tool list exercises exactly one ordering. + @aws.protocols#awsJson1_0 + service TestOneOfOutputRootService { + operations: [GetShape] + } + + @aws.protocols#awsJson1_0 + service TestOneOfInputRootService { + operations: [PutShape] + } + + /// Nested reference in the input (built first) caches the document's oneOf + /// schema; the output then requests the same shape as its root. + operation GetShape { + input: ShapeHolder + output: ShapeWithOneOf + } + + /// The polymorphic document is the input root (built first); the output then + /// references the same shape as a nested member. + operation PutShape { + input: ShapeWithOneOf + output: ShapeHolder + } + + @aws.protocols#awsJson1_0 + service TestUnionOutputRootService { + operations: [GetUnionShape] + } + + @aws.protocols#awsJson1_0 + service TestUnionInputRootService { + operations: [PutUnionShape] + } + + @aws.protocols#awsJson1_0 + service TestDocumentOutputRootService { + operations: [GetAnyDocument] + } + + /// Same orderings as GetShape/PutShape, with a plain union as the polymorphic root. + operation GetUnionShape { + input: UnionShapeHolder + output: ShapeUnion + } + + operation PutUnionShape { + input: ShapeUnion + output: UnionShapeHolder + } + + /// An untyped document as the output root. + operation GetAnyDocument { + output: AnyDocument + } + + structure ShapeHolder { + shape: ShapeWithOneOf + } + + structure UnionShapeHolder { + shape: ShapeUnion + } + + union ShapeUnion { + circle: Circle + square: Square + } + + document AnyDocument + + @oneOf(discriminator: "__type", members: [ + {name: "circle", target: Circle}, + {name: "square", target: Square} + ]) + document ShapeWithOneOf + + structure Circle { + @required + radius: Integer + } + + structure Square { + @required + side: Integer + }"""; + + // Assembled without validation, mirroring ModelBundles: bundled models reach the MCP server + // with document-typed operation inputs and outputs, which strict validation would reject. + private static final Model ONE_OF_ROOT_MODEL = Model.assembler() + .addUnparsedModel("test-oneof-root.smithy", ONE_OF_ROOT_MODEL_STR) + .discoverModels() + .disableValidation() + .assemble() + .unwrap(); + @Test void testUnionSchemaGeneratesOneOfWithWrappedMembers() { server = StdioMcpServer.builder() @@ -1902,6 +2006,122 @@ void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { assertEquals(2, oneOf.size(), "Document with @oneOf should have 2 oneOf variants"); } + private Map rootSchemaTools(String serviceName) { + server = StdioMcpServer.builder() + .name("smithy-mcp-server") + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test.oneofroot#" + serviceName)) + .proxyEndpoint("http://localhost") + .model(ONE_OF_ROOT_MODEL) + .build()) + .build(); + + server.start(); + + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); + write("tools/list", Document.of(Map.of())); + var response = read(); + var tools = new HashMap(); + for (var tool : response.getResult().asStringMap().get("tools").asList()) { + tools.put(tool.asStringMap().get("name").asString(), tool); + } + return tools; + } + + private static Map nestedShapeSchema(Document tool, String schemaKey) { + return tool.asStringMap() + .get(schemaKey) + .asStringMap() + .get("properties") + .asStringMap() + .get("shape") + .asStringMap(); + } + + @Test + void testOneOfDocumentAsOperationOutputRootWithCachedSchema() { + // The schema cache is per operation and the input is built first, so GetShape's nested + // input member caches the document's JsonOneOfSchema before the output requests the same + // shape as its root. This is the order that used to throw ClassCastException while + // building the tool list. + var tool = rootSchemaTools("TestOneOfOutputRootService").get("GetShape"); + + // The root must be object-typed (required by the MCP spec) and still carry the variants. + var outputSchema = tool.asStringMap().get("outputSchema").asStringMap(); + assertEquals("object", outputSchema.get("type").asString()); + assertEquals(2, outputSchema.get("oneOf").asList().size(), "Polymorphic output root should have 2 variants"); + + // The nested reference keeps its full oneOf schema. + assertEquals(2, + nestedShapeSchema(tool, "inputSchema").get("oneOf").asList().size(), + "Nested reference to the polymorphic document should keep its oneOf variants"); + } + + @Test + void testOneOfDocumentAsOperationInputRootBeforeNestedReference() { + // PutShape's input root is built first. This order used to render the member-less + // document as an empty object schema and cache it, so the nested output member then + // silently lost its oneOf variants. + var tool = rootSchemaTools("TestOneOfInputRootService").get("PutShape"); + + var inputSchema = tool.asStringMap().get("inputSchema").asStringMap(); + assertEquals("object", inputSchema.get("type").asString()); + assertEquals(2, inputSchema.get("oneOf").asList().size(), "Polymorphic input root should have 2 variants"); + + // Rendering the shape in an object position must not pollute the cache for nested uses. + assertEquals(2, + nestedShapeSchema(tool, "outputSchema").get("oneOf").asList().size(), + "Nested reference to the polymorphic document should keep its oneOf variants"); + } + + @Test + void testUnionAsOperationOutputRootWithCachedSchema() { + // Same ordering as GetShape with a plain union: the nested input member caches the + // union's JsonOneOfSchema before the output requests the union as its root. + var tool = rootSchemaTools("TestUnionOutputRootService").get("GetUnionShape"); + + var outputSchema = tool.asStringMap().get("outputSchema").asStringMap(); + assertEquals("object", outputSchema.get("type").asString()); + assertEquals(2, outputSchema.get("oneOf").asList().size(), "Union output root should have 2 variants"); + assertNull(outputSchema.get("properties"), "Union root must not render its members as properties"); + + assertEquals(2, + nestedShapeSchema(tool, "inputSchema").get("oneOf").asList().size(), + "Nested reference to the union should keep its oneOf variants"); + } + + @Test + void testUnionAsOperationInputRootBeforeNestedReference() { + // Same ordering as PutShape with a plain union: the union root is built first, then the + // output references it as a nested member. Rendering the root as a plain object would + // flatten the variants into sibling properties and cache that for the nested reference. + var tool = rootSchemaTools("TestUnionInputRootService").get("PutUnionShape"); + + var inputSchema = tool.asStringMap().get("inputSchema").asStringMap(); + assertEquals("object", inputSchema.get("type").asString()); + assertEquals(2, inputSchema.get("oneOf").asList().size(), "Union input root should have 2 variants"); + assertNull(inputSchema.get("properties"), "Union root must not render its members as properties"); + + assertEquals(2, + nestedShapeSchema(tool, "outputSchema").get("oneOf").asList().size(), + "Nested reference to the union should keep its oneOf variants"); + } + + @Test + void testPlainDocumentAsOperationOutputRoot() { + // An untyped document root has no members and no variants; it renders as a permissive + // object schema rather than failing or claiming an empty property set. + var tool = rootSchemaTools("TestDocumentOutputRootService").get("GetAnyDocument"); + + var outputSchema = tool.asStringMap().get("outputSchema").asStringMap(); + assertEquals("object", outputSchema.get("type").asString()); + assertNull(outputSchema.get("oneOf")); + assertNull(outputSchema.get("properties")); + } + @Test void testToolsListChangedNotificationInvalidatesCache() throws InterruptedException { var callCounter = new AtomicInteger(0);