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

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

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(

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);