Skip to content

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

Open
psp65 wants to merge 1 commit into
smithy-lang:mainfrom
psp65:mcp-oneof-root-schema
Open

fix(mcp-server): Handle oneOf documents as operation input/output roots#1345
psp65 wants to merge 1 commit into
smithy-lang:mainfrom
psp65:mcp-oneof-root-schema

Conversation

@psp65

@psp65 psp65 commented Sep 1, 2026

Copy link
Copy Markdown

What behavior changes?

McpService no longer fails (or silently degrades) when a shape carrying the smithy.mcp#oneOf trait is used directly as an operation input or output.

With a model like:

@oneOf(discriminator: "__type", members: [...])
document ShapeWithOneOf

operation GetShape {
    output: ShapeWithOneOf
}

operation ProcessShape {
    input: ProcessShapeInput // contains a member targeting ShapeWithOneOf
}

building the McpService previously had two order-dependent failures, because the per-service schema cache is shared across schema kinds:

  • If a nested reference (ProcessShape) was processed first, the cache held a JsonOneOfSchema for the shape, and createJsonObjectSchema for the GetShape output root then threw ClassCastException: JsonOneOfSchema cannot be cast to JsonObjectSchema — failing construction of the entire service (every tool), not just the offending operation.
  • If the root reference was processed first, createJsonObjectSchema treated the document as a plain (member-less) shape and cached an empty JsonObjectSchema under the shape id, which nested references then reused — silently dropping all oneOf variants.

After this change both positions render correctly: the root position produces {"type": "object", "oneOf": [...]} (the MCP spec requires tool input/output schemas to be object-typed; the oneOf variants are all objects, so the constraints compose), and nested references keep the cached JsonOneOfSchema exactly as before.

Why is this change needed?

Bundled models (loaded through ModelBundles, which assembles with validation disabled) can and do contain document-typed operation outputs: tooling that converts service models for MCP represents polymorphic type hierarchies as @oneOf documents, and a service with a polymorphic operation output currently cannot list tools at all. Observed in production as a deterministic ClassCastException at McpService.createJsonObjectSchema for every such service.

How was this validated?

  • Two new regression tests in McpServerTest, one per operation processing order (operations are processed in sorted order, so each order is pinned by operation naming):
    • testOneOfDocumentAsOperationOutputRoot — root processed first: asserts the object-position rendering of the @oneOf document root AND that a nested reference to the same shape keeps its variants (the silent cache-pollution mode).
    • testOneOfDocumentAsOperationOutputRootWithCachedSchema — nested reference processed first, so the root request hits a cached JsonOneOfSchema (the ClassCastException mode).
    • Both test models are assembled with disableValidation() to mirror ModelBundles.
  • Verified both tests fail on main without the fix and pass with it.
  • ./gradlew :mcp:mcp-server:test passes.

What should reviewers focus on?

  • McpService#asJsonObjectSchema: the re-shaping of a JsonOneOfSchema into an object-typed schema, and the new @oneOf guard in createJsonObjectSchema that routes trait-carrying shapes through createJsonOneOfSchema so the cache keeps the full oneOf schema for other references.
  • mcp-schemas/model/main.smithy: JsonObjectSchema gains an optional oneOf member so the object-typed root can carry the variants. The converted root preserves the type: "object", oneOf, and description a JsonOneOfSchema would serialize, plus the $schema annotation that all object-typed roots already carry.
  • The new guard in createJsonObjectSchema is scoped to ShapeType.DOCUMENT (the trait's selector), so a non-document shape incorrectly carrying the trait in an unvalidated model keeps its existing rendering — matching what runtime input/output adaptation recognizes.

Additional Links

None.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant