From 0aa46c1320795bda31cb607275e66e80ef8fae6f Mon Sep 17 00:00:00 2001 From: Param Parikh Date: Tue, 1 Sep 2026 18:53:32 +0000 Subject: [PATCH] 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-service schema cache is shared across schema kinds, so this failed in one of two order-dependent ways: - If a nested reference was processed first, the cache held a JsonOneOfSchema and createJsonObjectSchema for the root then threw ClassCastException, failing McpService construction entirely. - If the root was processed first, the document rendered as 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 createJsonOneOfSchema (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. Operations are processed in sorted order, so the two regression tests pin one processing order each; both fail on main (one with the original ClassCastException) and pass with the fix. --- mcp/mcp-schemas/model/main.smithy | 4 + .../smithy/java/mcp/server/McpService.java | 40 +++- .../smithy/java/mcp/server/McpServerTest.java | 172 ++++++++++++++++++ 3 files changed, 214 insertions(+), 2 deletions(-) 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);