Skip to content

Commit 7cc9fe1

Browse files
adwsinghadwsingh
committed
Send array of type in documents
1 parent 2d381ab commit 7cc9fe1

File tree

4 files changed

+136
-16
lines changed

4 files changed

+136
-16
lines changed

‎mcp/mcp-schemas/model/main.smithy‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,13 @@ structure JsonPrimitiveSchema {
149149
description: String
150150
}
151151

152+
structure JsonDocumentSchema {
153+
@required
154+
type: StringList
155+
156+
description: String
157+
}
158+
152159
enum JsonPrimitiveType {
153160
NUMBER = "number"
154161
STRING = "string"
@@ -159,7 +166,7 @@ enum JsonPrimitiveType {
159166
map PropertiesMap {
160167
key: String
161168

162-
/// one of JsonObjectSchema | JsonArraySchema | JsonPrimitiveSchema
169+
/// one of JsonObjectSchema | JsonArraySchema | JsonPrimitiveSchema | JsonDocumentSchema
163170
value: Document
164171
}
165172

‎mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java‎

Lines changed: 96 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,13 @@
2727
import java.util.List;
2828
import java.util.Map;
2929
import java.util.Set;
30+
import java.util.stream.Stream;
3031
import org.junit.jupiter.api.AfterEach;
3132
import org.junit.jupiter.api.BeforeEach;
3233
import org.junit.jupiter.api.Test;
34+
import org.junit.jupiter.params.ParameterizedTest;
35+
import org.junit.jupiter.params.provider.Arguments;
36+
import org.junit.jupiter.params.provider.MethodSource;
3337
import software.amazon.smithy.java.core.serde.document.Document;
3438
import software.amazon.smithy.java.io.ByteBufferUtils;
3539
import software.amazon.smithy.java.json.JsonCodec;
@@ -255,12 +259,18 @@ void testEchoSchemaTypeMapping() throws Exception {
255259
assertEquals("array", echoProps.path("integerList").path("type").asText());
256260
assertEquals("array", echoProps.path("nestedList").path("type").asText());
257261

258-
// Nested structure, map, union, and document should be objects
262+
// Nested structure, map, and union should be objects
259263
assertEquals("object", echoProps.path("nested").path("type").asText());
260264
assertEquals("object", echoProps.path("stringMap").path("type").asText());
261265
assertEquals("object", echoProps.path("nestedMap").path("type").asText());
262266
assertEquals("object", echoProps.path("unionValue").path("type").asText());
263-
assertEquals("object", echoProps.path("documentValue").path("type").asText());
267+
268+
// Document should have type as array of all JSON types
269+
var documentType = echoProps.path("documentValue").path("type");
270+
assertTrue(documentType.isArray(), "Document type should be an array");
271+
var types = new HashSet<String>();
272+
documentType.forEach(node -> types.add(node.asText()));
273+
assertEquals(Set.of("string", "number", "boolean", "object", "array", "null"), types);
264274
}
265275
}
266276

@@ -530,6 +540,90 @@ void testDocumentWithPrimitive() {
530540
assertEquals("just a string", echo.getMember("documentValue").asString());
531541
}
532542

543+
static Stream<Arguments> documentSchemaValidationTestCases() {
544+
return Stream.of(
545+
// String document
546+
Arguments.of("string document", Document.of("hello world")),
547+
// Number document (integer)
548+
Arguments.of("integer document", Document.of(42)),
549+
// Number document (double)
550+
Arguments.of("double document", Document.of(3.14159)),
551+
// Boolean document (true)
552+
Arguments.of("boolean true document", Document.of(true)),
553+
// Boolean document (false)
554+
Arguments.of("boolean false document", Document.of(false)),
555+
// Array document
556+
Arguments.of("array document",
557+
Document.of(List.of(
558+
Document.of("a"),
559+
Document.of(1),
560+
Document.of(true)))),
561+
// Object document
562+
Arguments.of("object document",
563+
Document.of(Map.of(
564+
"key",
565+
Document.of("value"),
566+
"number",
567+
Document.of(123)))),
568+
// Nested object document
569+
Arguments.of("nested object document",
570+
Document.of(Map.of(
571+
"outer",
572+
Document.of(Map.of(
573+
"inner",
574+
Document.of("deep value")))))),
575+
// Mixed array document
576+
Arguments.of("mixed array document",
577+
Document.of(List.of(
578+
Document.of("string"),
579+
Document.of(42),
580+
Document.of(true),
581+
Document.of(Map.of("nested", Document.of("value"))),
582+
Document.of(List.of(Document.of(1), Document.of(2)))))),
583+
// Empty object document
584+
Arguments.of("empty object document", Document.of(Map.of())));
585+
}
586+
587+
@ParameterizedTest(name = "{0}")
588+
@MethodSource("documentSchemaValidationTestCases")
589+
void testDocumentValidatesAgainstSchema(String description, Document documentValue) throws Exception {
590+
initializeLatestProtocol();
591+
592+
// Get output schema from raw JSON (using Jackson directly)
593+
write("tools/list", Document.of(Map.of()));
594+
var toolsResponseJson = readRawResponse();
595+
var toolsResponseNode = OBJECT_MAPPER.readTree(toolsResponseJson);
596+
var outputSchemaNode = toolsResponseNode
597+
.path("result")
598+
.path("tools")
599+
.get(0)
600+
.path("outputSchema");
601+
602+
// Create input with the document value
603+
var echoData = new HashMap<String, Document>();
604+
echoData.put("requiredField", Document.of("required"));
605+
echoData.put("documentValue", documentValue);
606+
607+
// Call tool and get raw JSON response (using Jackson directly)
608+
var params = Document.of(Map.of(
609+
"name",
610+
Document.of("McpEcho"),
611+
"arguments",
612+
createEchoInput(echoData)));
613+
write("tools/call", params);
614+
var callResponseJson = readRawResponse();
615+
var callResponseNode = OBJECT_MAPPER.readTree(callResponseJson);
616+
var structuredContentNode = callResponseNode.path("result").path("structuredContent");
617+
618+
assertNotNull(structuredContentNode, "Structured content should not be null for: " + description);
619+
620+
// Validate using Jackson-parsed JSON directly
621+
JsonSchema schema = SCHEMA_FACTORY.getSchema(outputSchemaNode);
622+
Set<ValidationMessage> errors = schema.validate(structuredContentNode);
623+
assertTrue(errors.isEmpty(),
624+
"Validation errors for " + description + ": " + errors);
625+
}
626+
533627
// ========== Enum Tests ==========
534628

535629
@Test

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import software.amazon.smithy.java.mcp.model.Capabilities;
3939
import software.amazon.smithy.java.mcp.model.InitializeResult;
4040
import software.amazon.smithy.java.mcp.model.JsonArraySchema;
41+
import software.amazon.smithy.java.mcp.model.JsonDocumentSchema;
4142
import software.amazon.smithy.java.mcp.model.JsonObjectSchema;
4243
import software.amazon.smithy.java.mcp.model.JsonPrimitiveSchema;
4344
import software.amazon.smithy.java.mcp.model.JsonPrimitiveType;
@@ -531,29 +532,28 @@ private static JsonObjectSchema createJsonObjectSchema(Schema schema, Set<ShapeI
531532

532533
var jsonSchema = switch (member.type()) {
533534
case LIST, SET -> createJsonArraySchema(member.memberTarget(), visited);
534-
case MAP, STRUCTURE, UNION, DOCUMENT -> createJsonObjectSchema(member.memberTarget(), visited);
535+
case MAP, STRUCTURE, UNION -> createJsonObjectSchema(member.memberTarget(), visited);
536+
case DOCUMENT -> createJsonDocumentSchema(member);
535537
default -> createJsonPrimitiveSchema(member);
536538
};
537539

538540
properties.put(name, Document.of(jsonSchema));
539541
}
540542

541543
visited.remove(targetId);
542-
var builder = JsonObjectSchema.builder()
544+
return JsonObjectSchema.builder()
543545
.properties(properties)
544546
.required(requiredProperties)
545-
.description(memberDescription(schema));
546-
if (type.isShapeType(ShapeType.DOCUMENT)) {
547-
builder.additionalProperties(true);
548-
}
549-
return builder.build();
547+
.description(memberDescription(schema))
548+
.build();
550549
}
551550

552551
private static JsonArraySchema createJsonArraySchema(Schema schema, Set<ShapeId> visited) {
553552
var listMember = schema.listMember();
554553
var items = switch (listMember.type()) {
555554
case LIST, SET -> createJsonArraySchema(listMember.memberTarget(), visited);
556-
case MAP, STRUCTURE, UNION, DOCUMENT -> createJsonObjectSchema(listMember.memberTarget(), visited);
555+
case MAP, STRUCTURE, UNION -> createJsonObjectSchema(listMember.memberTarget(), visited);
556+
case DOCUMENT -> createJsonDocumentSchema(listMember);
557557
default -> createJsonPrimitiveSchema(listMember);
558558
};
559559
return JsonArraySchema.builder()
@@ -577,6 +577,21 @@ private static JsonPrimitiveSchema createJsonPrimitiveSchema(Schema member) {
577577
.build();
578578
}
579579

580+
private static final List<String> DOCUMENT_TYPES = List.of(
581+
"string",
582+
"number",
583+
"boolean",
584+
"object",
585+
"array",
586+
"null");
587+
588+
private static JsonDocumentSchema createJsonDocumentSchema(Schema member) {
589+
return JsonDocumentSchema.builder()
590+
.type(DOCUMENT_TYPES)
591+
.description(memberDescription(member))
592+
.build();
593+
}
594+
580595
private static String memberDescription(Schema schema) {
581596
String description = null;
582597
var trait = schema.getTrait(TraitKey.DOCUMENTATION_TRAIT);

‎mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -984,11 +984,15 @@ private void validateNestedStructure(Map<String, Document> nestedSchema) {
984984
assertEquals("A string that's nested", nestedStr.get("description").asString());
985985

986986
var nestedDocument = properties.get("nestedDocument").asStringMap();
987-
assertEquals("object", nestedDocument.get("type").asString());
988-
assertEquals("http://json-schema.org/draft-07/schema#", nestedDocument.get("$schema").asString());
989-
assertTrue(nestedDocument.get("additionalProperties").asBoolean());
990-
assertTrue(nestedDocument.get("properties").asStringMap().isEmpty());
991-
assertTrue(nestedDocument.get("required").asList().isEmpty());
987+
var documentTypesDocs = nestedDocument.get("type").asList();
988+
var documentTypes = documentTypesDocs.stream().map(Document::asString).toList();
989+
assertEquals(6, documentTypes.size());
990+
assertTrue(documentTypes.contains("string"));
991+
assertTrue(documentTypes.contains("number"));
992+
assertTrue(documentTypes.contains("boolean"));
993+
assertTrue(documentTypes.contains("object"));
994+
assertTrue(documentTypes.contains("array"));
995+
assertTrue(documentTypes.contains("null"));
992996

993997
var recursive = properties.get("recursive").asStringMap();
994998
assertEquals("object", recursive.get("type").asString());

0 commit comments

Comments
 (0)