From 08f9b3a85edb410ad21629b165226829d579b01c Mon Sep 17 00:00:00 2001
From: Adwait Kumar Singh <adwsingh@amazon.com>
Date: Wed, 11 Mar 2026 12:54:14 -0700
Subject: [PATCH] Avoid O(N!) allocations for recusrive union variants

---
 .../mcp/server/McpServerIntegrationTest.java  | 18 ++++++++
 .../it/resources/META-INF/smithy/main.smithy  | 41 +++++++++++++++++++
 .../smithy/java/mcp/server/McpService.java    |  8 +++-
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java
index c8ed8a1e5..34394bb82 100644
--- a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java
+++ b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java
@@ -2182,6 +2182,24 @@ void testUnionMemberNamesInSchema() {
         assertEquals(3, memberNames.size(), "Should have exactly 3 members");
     }
 
+    @Test
+    void testRecursiveOneOfSchemaTerminatesWithoutInfiniteLoop() {
+        initializeLatestProtocol();
+        var schemas = getMcpEchoToolSchemas();
+        var inputSchemaNode = schemas.inputSchema().getSchemaNode();
+
+        // The recursiveTreeNode field should exist and have a oneOf array
+        var treeNodeSchema = inputSchemaNode.path("properties")
+                .path("echo")
+                .path("properties")
+                .path("recursiveTreeNode");
+        assertFalse(treeNodeSchema.isMissingNode(), "recursiveTreeNode should be in the schema");
+
+        var oneOf = treeNodeSchema.path("oneOf");
+        assertFalse(oneOf.isMissingNode(), "recursiveTreeNode should have oneOf");
+        assertEquals(10, oneOf.size(), "Recursive @oneOf should have 10 variants");
+    }
+
     // ========== Helper Methods ==========
 
     private void initializeWithProtocolVersion(ProtocolVersion protocolVersion) {
diff --git a/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy b/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy
index ebd6d3e6b..3bf77e861 100644
--- a/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy
+++ b/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy
@@ -170,6 +170,21 @@ structure Echo {
     // Helper to make CircleWithNested reachable for schema generation
     circleWithNested: CircleWithNested
 
+    // Recursive @oneOf document (for testing cycle detection in schema generation)
+    recursiveTreeNode: TreeNode
+
+    // Helpers to make recursive @oneOf targets reachable for schema generation
+    recNodeA: RecNodeA
+    recNodeB: RecNodeB
+    recNodeC: RecNodeC
+    recNodeD: RecNodeD
+    recNodeE: RecNodeE
+    recNodeF: RecNodeF
+    recNodeG: RecNodeG
+    recNodeH: RecNodeH
+    recNodeI: RecNodeI
+    recNodeJ: RecNodeJ
+
     // Required field to test required validation
     @required
     requiredField: String
@@ -255,3 +270,29 @@ map ShapeWithOneOfMap {
     key: String
     value: ShapeWithOneOf
 }
+
+/// A recursive @oneOf document where multiple members cycle back, causing O(N!) schema blowup
+@oneOf(discriminator: "__type", members: [
+    {name: "nodeA", target: RecNodeA},
+    {name: "nodeB", target: RecNodeB},
+    {name: "nodeC", target: RecNodeC},
+    {name: "nodeD", target: RecNodeD},
+    {name: "nodeE", target: RecNodeE},
+    {name: "nodeF", target: RecNodeF},
+    {name: "nodeG", target: RecNodeG},
+    {name: "nodeH", target: RecNodeH},
+    {name: "nodeI", target: RecNodeI},
+    {name: "nodeJ", target: RecNodeJ}
+])
+document TreeNode
+
+structure RecNodeA { value: String, child: TreeNode }
+structure RecNodeB { value: String, child: TreeNode }
+structure RecNodeC { value: String, child: TreeNode }
+structure RecNodeD { value: String, child: TreeNode }
+structure RecNodeE { value: String, child: TreeNode }
+structure RecNodeF { value: String, child: TreeNode }
+structure RecNodeG { value: String, child: TreeNode }
+structure RecNodeH { value: String, child: TreeNode }
+structure RecNodeI { value: String, child: TreeNode }
+structure RecNodeJ { value: String, child: TreeNode }
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 a184ed0fc..7fb7bae06 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
@@ -694,11 +694,16 @@ private SerializableShape createJsonDocumentSchema(Schema member, Set<ShapeId> v
         }
     }
 
-    private JsonOneOfSchema createJsonOneOfSchema(
+    private SerializableShape createJsonOneOfSchema(
             OneOfTrait oneOfTrait,
             Schema documentMember,
             Set<ShapeId> visited
     ) {
+        var targetId = (documentMember.isMember() ? documentMember.memberTarget() : documentMember).id();
+        if (!visited.add(targetId)) {
+            return JsonObjectSchema.builder().build();
+        }
+
         var oneOfVariants = new ArrayList<Document>();
 
         for (var memberDef : oneOfTrait.getMembers()) {
@@ -711,6 +716,7 @@ private JsonOneOfSchema createJsonOneOfSchema(
             oneOfVariants.add(createUnionVariant(memberName, memberSchema));
         }
 
+        visited.remove(targetId);
         return JsonOneOfSchema.builder()
                 .oneOf(oneOfVariants)
                 .description(memberDescription(documentMember))