diff --git a/mcp/mcp-schemas/model/main.smithy b/mcp/mcp-schemas/model/main.smithy index 047326123..12fc8c2a7 100644 --- a/mcp/mcp-schemas/model/main.smithy +++ b/mcp/mcp-schemas/model/main.smithy @@ -125,6 +125,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/McpService.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java index a0be51256..e633e6354 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java @@ -992,7 +992,21 @@ private JsonObjectSchema createJsonObjectSchema( 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 oneOfTrait = target.getTrait(ONE_OF_TRAIT); + if (oneOfTrait != null) { + return asJsonObjectSchema(createJsonOneOfSchema(oneOfTrait, member, visited, cache)); + } } if (!visited.add(targetId)) { @@ -1022,7 +1036,29 @@ private JsonObjectSchema createJsonObjectSchema( .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, which serializes to the + * same JSON. Anything else degrades to a permissive object schema rather than failing the + * entire tool listing. + */ + private static JsonObjectSchema asJsonObjectSchema(SerializableShape schema) { + if (schema instanceof JsonObjectSchema objectSchema) { + return objectSchema; + } + if (schema instanceof JsonOneOfSchema oneOfSchema) { + var builder = JsonObjectSchema.builder().oneOf(oneOfSchema.getOneOf()); + if (oneOfSchema.getDescription() != null) { + builder.description(oneOfSchema.getDescription()); + } + return builder.build(); + } + return JsonObjectSchema.builder().build(); } private JsonArraySchema createJsonArraySchema( diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java index 3d8c2df25..6d2dd88d3 100644 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java @@ -1580,6 +1580,87 @@ 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 + + @aws.protocols#awsJson1_0 + service TestOneOfRootService { + operations: [ProcessShape, GetShape] + } + + /// References the polymorphic document as a nested member, so the schema + /// cache holds a oneOf schema for it when GetShape's output is built. + operation ProcessShape { + input: ProcessShapeInput + output: ProcessShapeOutput + } + + /// Returns the polymorphic document directly as the operation output. + operation GetShape { + input: GetShapeInput + output: ShapeWithOneOf + } + + structure ProcessShapeInput { + shape: ShapeWithOneOf + } + + structure ProcessShapeOutput { + shape: ShapeWithOneOf + } + + structure GetShapeInput { + id: String + } + + /// Same shapes, but operation names sort the nested reference BEFORE the + /// root reference: AProcessShape caches the oneOf schema, then ZGetShape + /// requests the same shape as its output root (the cache-hit order). + @aws.protocols#awsJson1_0 + service TestOneOfRootCachedService { + operations: [AProcessShape, ZGetShape] + } + + operation AProcessShape { + input: ProcessShapeInput + output: ProcessShapeOutput + } + + operation ZGetShape { + input: GetShapeInput + output: 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 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 = McpServer.builder() @@ -1672,6 +1753,97 @@ void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { assertEquals(2, oneOf.size(), "Document with @oneOf should have 2 oneOf variants"); } + @Test + void testOneOfDocumentAsOperationOutputRoot() { + // Regression test: a @oneOf document used directly as an operation output used to + // either throw ClassCastException (when another operation had already cached its + // JsonOneOfSchema) or silently cache an empty object schema that clobbered nested + // references, depending on operation processing order. + server = McpServer.builder() + .name("smithy-mcp-server") + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test.oneofroot#TestOneOfRootService")) + .proxyEndpoint("http://localhost") + .model(ONE_OF_ROOT_MODEL) + .build()) + .build(); + + server.start(); + + initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + write("tools/list", Document.of(Map.of())); + var response = read(); + var tools = response.getResult().asStringMap().get("tools").asList(); + + // The tool whose output IS the polymorphic document: the root schema must be + // object-typed (required by the MCP spec) and still carry the oneOf variants. + var rootTool = tools.stream() + .filter(t -> t.asStringMap().get("name").asString().equals("GetShape")) + .findFirst() + .orElseThrow() + .asStringMap(); + var outputSchema = rootTool.get("outputSchema").asStringMap(); + assertEquals("object", outputSchema.get("type").asString()); + var rootOneOf = outputSchema.get("oneOf").asList(); + assertEquals(2, rootOneOf.size(), "Polymorphic output root should have 2 oneOf variants"); + + // The nested reference must keep its full oneOf schema: rendering the same shape in an + // object position must not pollute the schema cache for other references. + var nestedTool = tools.stream() + .filter(t -> t.asStringMap().get("name").asString().equals("ProcessShape")) + .findFirst() + .orElseThrow() + .asStringMap(); + var shapeProp = nestedTool.get("inputSchema") + .asStringMap() + .get("properties") + .asStringMap() + .get("shape") + .asStringMap(); + assertEquals(2, + shapeProp.get("oneOf").asList().size(), + "Nested reference to the polymorphic document should keep its oneOf variants"); + } + + @Test + void testOneOfDocumentAsOperationOutputRootWithCachedSchema() { + // Operations are processed in sorted order, so AProcessShape caches the document's + // JsonOneOfSchema before ZGetShape requests the same shape as its output root. This is + // the order that used to throw ClassCastException while building the service. + server = McpServer.builder() + .name("smithy-mcp-server") + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test.oneofroot#TestOneOfRootCachedService")) + .proxyEndpoint("http://localhost") + .model(ONE_OF_ROOT_MODEL) + .build()) + .build(); + + server.start(); + + initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + write("tools/list", Document.of(Map.of())); + var response = read(); + var tools = response.getResult().asStringMap().get("tools").asList(); + + var rootTool = tools.stream() + .filter(t -> t.asStringMap().get("name").asString().equals("ZGetShape")) + .findFirst() + .orElseThrow() + .asStringMap(); + var outputSchema = rootTool.get("outputSchema").asStringMap(); + assertEquals("object", outputSchema.get("type").asString()); + assertEquals(2, + outputSchema.get("oneOf").asList().size(), + "Root schema converted from the cached oneOf schema should keep its variants"); + } + @Test void testToolsListChangedNotificationInvalidatesCache() throws InterruptedException { var callCounter = new AtomicInteger(0);