fix(mcp-server): Handle oneOf documents as operation input/output roots

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#"
 }

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<ShapeId, SerializableShape> 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<ShapeId> visited,
             Map<ShapeId, SerializableShape> 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<Document>();
         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<ShapeId> visited,
             Map<ShapeId, SerializableShape> 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<ShapeId> visited,
             Map<ShapeId, SerializableShape> 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",

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<String, Document> 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<String, Document>();
+        for (var tool : response.getResult().asStringMap().get("tools").asList()) {
+            tools.put(tool.asStringMap().get("name").asString(), tool);
+        }
+        return tools;
+    }
+
+    private static Map<String, Document> 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);