From dda9ec9050a6cd2bcf8f65075fe07551c7fac31e Mon Sep 17 00:00:00 2001 From: Param Parikh Date: Mon, 14 Sep 2026 17:45:33 +0000 Subject: [PATCH 1/2] fix(mcp-server): Handle oneOf documents as operation input/output roots A document carrying the smithy.mcp#oneOf trait (a discriminated polymorphic type) can appear as an operation's input or output in bundled models, which are loaded with validation disabled. The per-operation schema cache is shared across schema kinds and the input schema is built before the output, so this failed in one of two order-dependent ways: - If a nested reference was rendered first, the cache held a JsonOneOfSchema and createObjectSchema for the root then threw ClassCastException, failing tool listing for the whole service. - If the root was rendered first, the document produced an empty object schema that was cached under the shape id, silently dropping the oneOf variants from every nested reference. Route oneOf documents in object positions through createOneOfSchema (preserving the cached oneOf schema for other references) and re-shape the result into an object-typed schema: JsonObjectSchema gains an optional oneOf member, producing {"type": "object", "oneOf": [...]} as the MCP spec requires for tool schemas. The guard is scoped to ShapeType.DOCUMENT (the trait's selector), so any other shape kind carrying the trait keeps its regular rendering, matching what runtime input/output adaptation recognizes. Two regression tests pin one ordering each within a single operation (nested reference then document root, and document root then nested reference); both fail on main and each fails only for its own mechanism when that mechanism is mutated out. --- mcp/mcp-schemas/model/main.smithy | 4 + .../java/mcp/server/McpSchemaFactory.java | 40 +++++- .../java/mcp/server/StdioMcpServerTest.java | 133 ++++++++++++++++++ 3 files changed, 175 insertions(+), 2 deletions(-) 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..b2ed97703 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 @@ -126,8 +126,23 @@ private JsonObjectSchema createObjectSchema( var targetId = target.id(); var cached = cache.get(targetId); if (cached != null) { - return (JsonObjectSchema) withDescription(cached, memberDescription(member)); + return asJsonObjectSchema(withDescription(cached, memberDescription(member))); } + + // A document carrying the oneOf trait (a discriminated polymorphic type) can be asked + // for in an object position — most notably as an operation's input or output, which + // model bundles load without validation. Build it through the oneOf path, which caches + // a JsonOneOfSchema for other references to reuse, and re-shape the result into the + // object-typed schema this position requires. Scoped to documents (the trait's + // selector) so any other shape kind carrying the trait keeps its regular rendering, + // matching what runtime input/output adaptation recognizes. + if (target.type() == ShapeType.DOCUMENT) { + var oneOf = target.getTrait(ONE_OF_TRAIT); + if (oneOf != null) { + return asJsonObjectSchema(createOneOfSchema(oneOf, member, visited, cache)); + } + } + if (!visited.add(targetId)) { return JsonObjectSchema.builder().build(); } @@ -147,7 +162,28 @@ private JsonObjectSchema createObjectSchema( .required(required) .build(); cache.put(targetId, result); - return (JsonObjectSchema) withDescription(result, memberDescription(member)); + return asJsonObjectSchema(withDescription(result, memberDescription(member))); + } + + /** + * Re-shapes a schema for a position that requires an object-typed schema, such as a tool's + * input or output (the MCP spec requires both to have {@code "type": "object"}). A + * discriminated polymorphic type renders as a {@link JsonOneOfSchema}; it is carried over as + * an object schema constrained by the same {@code oneOf} variants. Anything else degrades to + * a permissive object schema rather than failing the entire tool listing. + */ + private static JsonObjectSchema asJsonObjectSchema(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(); + } + default -> JsonObjectSchema.builder().build(); + }; } private JsonArraySchema createArraySchema( 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..f0bec6051 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,68 @@ 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 + } + + structure ShapeHolder { + shape: ShapeWithOneOf + } + + @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 +1964,77 @@ void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { assertEquals(2, oneOf.size(), "Document with @oneOf should have 2 oneOf variants"); } + private Map oneOfRootTools(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 = oneOfRootTools("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 = oneOfRootTools("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 testToolsListChangedNotificationInvalidatesCache() throws InterruptedException { var callCounter = new AtomicInteger(0); From a1c7a1584dfa620da0409e3449fa9ff86d867604 Mon Sep 17 00:00:00 2001 From: Adwait Kumar Singh Date: Tue, 15 Sep 2026 04:17:39 +0530 Subject: [PATCH 2/2] Adapt root input/output schemas also --- .../java/mcp/server/McpSchemaFactory.java | 88 ++++++++---------- .../java/mcp/server/StdioMcpServerTest.java | 93 ++++++++++++++++++- 2 files changed, 130 insertions(+), 51 deletions(-) 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 b2ed97703..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( @@ -126,23 +118,8 @@ private JsonObjectSchema createObjectSchema( var targetId = target.id(); var cached = cache.get(targetId); if (cached != null) { - return asJsonObjectSchema(withDescription(cached, memberDescription(member))); + return (JsonObjectSchema) withDescription(cached, memberDescription(member)); } - - // A document carrying the oneOf trait (a discriminated polymorphic type) can be asked - // for in an object position — most notably as an operation's input or output, which - // model bundles load without validation. Build it through the oneOf path, which caches - // a JsonOneOfSchema for other references to reuse, and re-shape the result into the - // object-typed schema this position requires. Scoped to documents (the trait's - // selector) so any other shape kind carrying the trait keeps its regular rendering, - // matching what runtime input/output adaptation recognizes. - if (target.type() == ShapeType.DOCUMENT) { - var oneOf = target.getTrait(ONE_OF_TRAIT); - if (oneOf != null) { - return asJsonObjectSchema(createOneOfSchema(oneOf, member, visited, cache)); - } - } - if (!visited.add(targetId)) { return JsonObjectSchema.builder().build(); } @@ -153,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); @@ -162,17 +139,22 @@ private JsonObjectSchema createObjectSchema( .required(required) .build(); cache.put(targetId, result); - return asJsonObjectSchema(withDescription(result, memberDescription(member))); + return (JsonObjectSchema) withDescription(result, memberDescription(member)); + } + + private JsonObjectSchema createRootSchema(Schema root, Map cache) { + return asJsonObjectSchema(root, createSchema(root, new HashSet<>(), cache)); } /** - * Re-shapes a schema for a position that requires an object-typed schema, such as a tool's - * input or output (the MCP spec requires both to have {@code "type": "object"}). A - * discriminated polymorphic type renders as a {@link JsonOneOfSchema}; it is carried over as - * an object schema constrained by the same {@code oneOf} variants. Anything else degrades to - * a permissive object schema rather than failing the entire tool listing. + * 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(SerializableShape schema) { + private static JsonObjectSchema asJsonObjectSchema(Schema root, SerializableShape schema) { return switch (schema) { case JsonObjectSchema object -> object; case JsonOneOfSchema oneOf -> { @@ -182,6 +164,13 @@ private static JsonObjectSchema asJsonObjectSchema(SerializableShape schema) { } 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(); }; } @@ -192,7 +181,7 @@ private JsonArraySchema createArraySchema( 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", @@ -269,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); @@ -300,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); @@ -309,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); }; } @@ -330,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 f0bec6051..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 @@ -1843,10 +1843,52 @@ private void writeNotification(String method, Document params) { 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} @@ -1964,7 +2006,7 @@ void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { assertEquals(2, oneOf.size(), "Document with @oneOf should have 2 oneOf variants"); } - private Map oneOfRootTools(String serviceName) { + private Map rootSchemaTools(String serviceName) { server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) @@ -2005,7 +2047,7 @@ void testOneOfDocumentAsOperationOutputRootWithCachedSchema() { // 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 = oneOfRootTools("TestOneOfOutputRootService").get("GetShape"); + 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(); @@ -2023,7 +2065,7 @@ 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 = oneOfRootTools("TestOneOfInputRootService").get("PutShape"); + var tool = rootSchemaTools("TestOneOfInputRootService").get("PutShape"); var inputSchema = tool.asStringMap().get("inputSchema").asStringMap(); assertEquals("object", inputSchema.get("type").asString()); @@ -2035,6 +2077,51 @@ void testOneOfDocumentAsOperationInputRootBeforeNestedReference() { "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);