Skip to content

Commit b773ad5

Browse files
jknackjknack
committed
feat(mcp): introduce @McpOutputSchema and runtime generation controls
To conserve LLM context window tokens and avoid recursive reflection on complex return types, this commit makes JSON output schema generation for MCP tools configurable and strictly opt-in. * Adds a stateful `generateOutputSchema` flag to the APT-generated `McpService` classes (defaults to false). This allows developers to toggle global output schema generation at runtime via the `McpModule`. * Introduces the top-level `@McpOutputSchema` annotation for granular, per-tool overrides, independent of the global flag. - `@McpOutputSchema.Off`: Explicitly disables schema generation. - `@McpOutputSchema.From(Class)`: Forces generation using a specific type, elegantly bypassing Java type erasure (e.g., for generic Maps). - `@McpOutputSchema.ArrayOf(Class)`: Forces generation as a JSON array. - `@McpOutputSchema.MapOf(Class)`: Forces generation as a JSON object map. * Updates `McpRouter` to inject the stateful flag and setter method into generated routing classes. * Updates `McpRoute` to conditionally generate Victools schema extraction logic based on annotation presence and the runtime fallback flag.
1 parent f722ce8 commit b773ad5

File tree

10 files changed

+494
-24
lines changed

10 files changed

+494
-24
lines changed

‎modules/jooby-apt/src/main/java/io/jooby/internal/apt/McpRoute.java‎

Lines changed: 175 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -492,7 +492,6 @@ private List<String> generateToolDefinition(boolean kt) {
492492
")"));
493493
}
494494

495-
// Switched from .set() to .put() for standard Map
496495
buffer.add(statement(indent(6), "props.put(", string(mcpName), ", schema_", mcpName, ")"));
497496

498497
if (!param.isNullable(kt)) {
@@ -523,7 +522,6 @@ private List<String> generateToolDefinition(boolean kt) {
523522
semicolon(kt)));
524523
}
525524

526-
// Switched from .set() to .put() for standard Map
527525
buffer.add(
528526
statement(
529527
indent(6),
@@ -540,52 +538,157 @@ private List<String> generateToolDefinition(boolean kt) {
540538
}
541539
}
542540

543-
// --- OUTPUT SCHEMA GENERATION ---
544-
String returnTypeStr = getReturnType().getRawType().toString();
545-
boolean generateOutputSchema = hasOutputSchema();
541+
// --- OUTPUT SCHEMA GENERATION (RUNTIME AWARE) ---
542+
var outMeta = parseOutputSchemaMeta();
543+
boolean isEligible = hasOutputSchema();
544+
546545
String outputSchemaArg = "null";
547546

548-
if (generateOutputSchema) {
547+
if (outMeta.isOff() || (!isEligible && outMeta.type() == null)) {
548+
// Do nothing, outputSchemaArg remains "null"
549+
} else {
549550
outputSchemaArg = getMethodName() + "OutputSchema";
551+
String targetTypeStr =
552+
outMeta.type() != null ? outMeta.type() : getReturnType().getRawType().toString();
553+
550554
if (kt) {
555+
buffer.add(
556+
statement(indent(6), "var ", outputSchemaArg, ": java.util.Map<String, Any>? = null"));
557+
} else {
551558
buffer.add(
552559
statement(
553560
indent(6),
561+
"java.util.Map<String, Object> ",
562+
outputSchemaArg,
563+
" = null",
564+
semicolon(kt)));
565+
}
566+
567+
boolean needsRuntimeCheck = (outMeta.type() == null);
568+
String ind = indent(6);
569+
570+
if (needsRuntimeCheck) {
571+
buffer.add(statement(indent(6), "if (this.generateOutputSchema) {"));
572+
ind = indent(8);
573+
}
574+
575+
if (kt) {
576+
buffer.add(
577+
statement(
578+
ind,
554579
"val ",
555580
outputSchemaArg,
556581
"Node = schemaGenerator.generateSchema(",
557-
returnTypeStr,
582+
targetTypeStr,
558583
"::class.java)"));
559-
// Use this.json to convert the output schema
560584
buffer.add(
561585
statement(
562-
indent(6),
586+
ind,
563587
"val ",
564588
outputSchemaArg,
565-
" = this.json.convertValue(",
589+
"Map = this.json.convertValue(",
566590
outputSchemaArg,
567591
"Node, java.util.Map::class.java) as java.util.Map<String, Any>"));
568592
} else {
569593
buffer.add(
570594
statement(
571-
indent(6),
595+
ind,
572596
"var ",
573597
outputSchemaArg,
574598
"Node = schemaGenerator.generateSchema(",
575-
returnTypeStr,
599+
targetTypeStr,
576600
".class)",
577601
semicolon(kt)));
578-
// Use this.json to convert the output schema
579602
buffer.add(
580603
statement(
581-
indent(6),
604+
ind,
582605
"var ",
583606
outputSchemaArg,
584-
" = this.json.convertValue(",
607+
"Map = this.json.convertValue(",
585608
outputSchemaArg,
586609
"Node, java.util.Map.class)",
587610
semicolon(kt)));
588611
}
612+
613+
// Handle ArrayOf and MapOf Schema Wrapping
614+
if (outMeta.schemaType() == SchemaType.ARRAY) {
615+
if (kt) {
616+
buffer.add(
617+
statement(
618+
ind,
619+
"val ",
620+
outputSchemaArg,
621+
"Wrapped = java.util.LinkedHashMap<String, Any>()"));
622+
buffer.add(statement(ind, outputSchemaArg, "Wrapped.put(\"type\", \"array\")"));
623+
buffer.add(
624+
statement(ind, outputSchemaArg, "Wrapped.put(\"items\", ", outputSchemaArg, "Map)"));
625+
buffer.add(statement(ind, outputSchemaArg, " = ", outputSchemaArg, "Wrapped"));
626+
} else {
627+
buffer.add(
628+
statement(
629+
ind,
630+
"var ",
631+
outputSchemaArg,
632+
"Wrapped = new java.util.LinkedHashMap<String, Object>()",
633+
semicolon(kt)));
634+
buffer.add(
635+
statement(ind, outputSchemaArg, "Wrapped.put(\"type\", \"array\")", semicolon(kt)));
636+
buffer.add(
637+
statement(
638+
ind,
639+
outputSchemaArg,
640+
"Wrapped.put(\"items\", ",
641+
outputSchemaArg,
642+
"Map)",
643+
semicolon(kt)));
644+
buffer.add(
645+
statement(ind, outputSchemaArg, " = ", outputSchemaArg, "Wrapped", semicolon(kt)));
646+
}
647+
} else if (outMeta.schemaType() == SchemaType.MAP) {
648+
if (kt) {
649+
buffer.add(
650+
statement(
651+
ind,
652+
"val ",
653+
outputSchemaArg,
654+
"Wrapped = java.util.LinkedHashMap<String, Any>()"));
655+
buffer.add(statement(ind, outputSchemaArg, "Wrapped.put(\"type\", \"object\")"));
656+
buffer.add(
657+
statement(
658+
ind,
659+
outputSchemaArg,
660+
"Wrapped.put(\"additionalProperties\", ",
661+
outputSchemaArg,
662+
"Map)"));
663+
buffer.add(statement(ind, outputSchemaArg, " = ", outputSchemaArg, "Wrapped"));
664+
} else {
665+
buffer.add(
666+
statement(
667+
ind,
668+
"var ",
669+
outputSchemaArg,
670+
"Wrapped = new java.util.LinkedHashMap<String, Object>()",
671+
semicolon(kt)));
672+
buffer.add(
673+
statement(ind, outputSchemaArg, "Wrapped.put(\"type\", \"object\")", semicolon(kt)));
674+
buffer.add(
675+
statement(
676+
ind,
677+
outputSchemaArg,
678+
"Wrapped.put(\"additionalProperties\", ",
679+
outputSchemaArg,
680+
"Map)",
681+
semicolon(kt)));
682+
buffer.add(
683+
statement(ind, outputSchemaArg, " = ", outputSchemaArg, "Wrapped", semicolon(kt)));
684+
}
685+
} else {
686+
buffer.add(statement(ind, outputSchemaArg, " = ", outputSchemaArg, "Map", semicolon(kt)));
687+
}
688+
689+
if (needsRuntimeCheck) {
690+
buffer.add(statement(indent(6), "}"));
691+
}
589692
}
590693

591694
// --- NESTED ANNOTATION EXTRACTION ---
@@ -622,7 +725,6 @@ private List<String> generateToolDefinition(boolean kt) {
622725
titleArg,
623726
", ",
624727
string(description),
625-
// Use this.json to convert the main schema map into JsonSchema
626728
", this.json.convertValue(schema,"
627729
+ " io.modelcontextprotocol.spec.McpSchema.JsonSchema::class.java), ",
628730
outputSchemaArg,
@@ -660,7 +762,6 @@ private List<String> generateToolDefinition(boolean kt) {
660762
titleArg,
661763
", ",
662764
string(description),
663-
// Use this.json to convert the main schema map into JsonSchema
664765
", this.json.convertValue(schema,"
665766
+ " io.modelcontextprotocol.spec.McpSchema.JsonSchema.class), ",
666767
outputSchemaArg,
@@ -708,7 +809,7 @@ public List<String> generateMcpHandlerMethod(boolean kt) {
708809
"private fun ",
709810
handlerName,
710811
"(exchange: io.modelcontextprotocol.server.McpSyncServerExchange?, transportContext:"
711-
+ " io.modelcontextprotocol.common.McpTransportContext, req:" // Removed '?'
812+
+ " io.modelcontextprotocol.common.McpTransportContext, req:"
712813
+ " io.modelcontextprotocol.spec.McpSchema.",
713814
reqType,
714815
"): io.modelcontextprotocol.spec.McpSchema.",
@@ -727,7 +828,7 @@ public List<String> generateMcpHandlerMethod(boolean kt) {
727828
handlerName,
728829
"(io.modelcontextprotocol.server.McpSyncServerExchange exchange,"
729830
+ " io.modelcontextprotocol.common.McpTransportContext"
730-
+ " transportContext," // Guaranteed non-null
831+
+ " transportContext,"
731832
+ " io.modelcontextprotocol.spec.McpSchema.",
732833
reqType,
733834
" req) {"));
@@ -969,11 +1070,22 @@ public List<String> generateMcpHandlerMethod(boolean kt) {
9691070

9701071
var methodCall = "c." + getMethodName() + "(" + String.join(", ", javaParamNames) + ")";
9711072

972-
// Prefix for Resources: "req.uri(), "
9731073
String toMethodPrefix = (isMcpResource() || isMcpResourceTemplate()) ? "req.uri(), " : "";
9741074

975-
// Suffix for Tools: ", true" or ", false"
976-
String toMethodSuffix = isMcpTool() ? ", " + hasOutputSchema() : "";
1075+
// Resolve output schema flag for Handler runtime behavior
1076+
String toMethodSuffix = "";
1077+
if (isMcpTool()) {
1078+
var outMeta = parseOutputSchemaMeta();
1079+
boolean isEligible = hasOutputSchema();
1080+
1081+
if (outMeta.isOff() || (!isEligible && outMeta.type() == null)) {
1082+
toMethodSuffix = ", false";
1083+
} else if (outMeta.type() != null) {
1084+
toMethodSuffix = ", true";
1085+
} else {
1086+
toMethodSuffix = ", this.generateOutputSchema";
1087+
}
1088+
}
9771089

9781090
if (getReturnType().isVoid()) {
9791091
buffer.add(statement(indent(6), methodCall, semicolon(kt)));
@@ -1070,6 +1182,38 @@ private boolean hasOutputSchema() {
10701182
&& !isMcpClass;
10711183
}
10721184

1185+
private String extractClassValue(String annotationName) {
1186+
var annotation = AnnotationSupport.findAnnotationByName(method, annotationName);
1187+
if (annotation == null) return null;
1188+
var val =
1189+
AnnotationSupport.findAnnotationValue(annotation, "value"::equals).stream()
1190+
.findFirst()
1191+
.orElse(null);
1192+
if (val != null && val.endsWith(".class")) {
1193+
return val.substring(0, val.length() - 6);
1194+
}
1195+
return val;
1196+
}
1197+
1198+
private OutputSchemaMeta parseOutputSchemaMeta() {
1199+
if (AnnotationSupport.findAnnotationByName(
1200+
method, "io.jooby.annotation.mcp.McpOutputSchema.Off")
1201+
!= null) {
1202+
return new OutputSchemaMeta(true, null, SchemaType.NONE);
1203+
}
1204+
1205+
String fromType = extractClassValue("io.jooby.annotation.mcp.McpOutputSchema.From");
1206+
if (fromType != null) return new OutputSchemaMeta(false, fromType, SchemaType.FROM);
1207+
1208+
String arrayType = extractClassValue("io.jooby.annotation.mcp.McpOutputSchema.ArrayOf");
1209+
if (arrayType != null) return new OutputSchemaMeta(false, arrayType, SchemaType.ARRAY);
1210+
1211+
String mapType = extractClassValue("io.jooby.annotation.mcp.McpOutputSchema.MapOf");
1212+
if (mapType != null) return new OutputSchemaMeta(false, mapType, SchemaType.MAP);
1213+
1214+
return new OutputSchemaMeta(false, null, SchemaType.NONE);
1215+
}
1216+
10731217
private McpAnnotation parseResourceAnnotation() {
10741218
String rawAnnotations =
10751219
extractAnnotationValue("io.jooby.annotation.mcp.McpResource", "annotations");
@@ -1138,4 +1282,13 @@ private record McpToolAnnotation(
11381282
String readOnlyHint, String destructiveHint, String idempotentHint, String openWorldHint) {}
11391283

11401284
private record McpAnnotation(String audience, String lastModified, String priority) {}
1285+
1286+
private enum SchemaType {
1287+
NONE,
1288+
FROM,
1289+
ARRAY,
1290+
MAP
1291+
}
1292+
1293+
private record OutputSchemaMeta(boolean isOff, String type, SchemaType schemaType) {}
11411294
}

‎modules/jooby-apt/src/main/java/io/jooby/internal/apt/McpRouter.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,10 +188,13 @@ public String toSourceCode(boolean kt) throws IOException {
188188
buffer.append(
189189
statement(
190190
indent(4), "private lateinit var json: io.modelcontextprotocol.json.McpJsonMapper"));
191+
buffer.append(statement(indent(4), "private var generateOutputSchema: Boolean = false"));
191192
} else {
192193
buffer.append(
193194
statement(
194195
indent(4), "private io.modelcontextprotocol.json.McpJsonMapper json", semicolon(kt)));
196+
buffer.append(
197+
statement(indent(4), "private boolean generateOutputSchema = false", semicolon(kt)));
195198
}
196199

197200
// --- capabilities() ---
@@ -224,6 +227,29 @@ public String toSourceCode(boolean kt) throws IOException {
224227
}
225228
buffer.append(statement(indent(4), "}\n"));
226229

230+
// --- generateOutputSchema() ---
231+
if (kt) {
232+
buffer.append(
233+
statement(
234+
indent(4),
235+
"override fun generateOutputSchema(generateOutputSchema: Boolean):"
236+
+ " io.jooby.mcp.McpService {"));
237+
buffer.append(statement(indent(6), "this.generateOutputSchema = generateOutputSchema"));
238+
buffer.append(statement(indent(6), "return this"));
239+
buffer.append(statement(indent(4), "}\n"));
240+
} else {
241+
buffer.append(statement(indent(4), "@Override"));
242+
buffer.append(
243+
statement(
244+
indent(4),
245+
"public io.jooby.mcp.McpService generateOutputSchema(boolean generateOutputSchema)"
246+
+ " {"));
247+
buffer.append(
248+
statement(indent(6), "this.generateOutputSchema = generateOutputSchema", semicolon(kt)));
249+
buffer.append(statement(indent(6), "return this", semicolon(kt)));
250+
buffer.append(statement(indent(4), "}\n"));
251+
}
252+
227253
// --- serverKey() ---
228254
var serverName = getMcpServerKey();
229255
if (kt) {

‎modules/jooby-apt/src/test/java/tests/i3830/Issue3830.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ private void setup(java.util.function.Function<io.jooby.Context, ExampleServer>
4747
}
4848
4949
private io.modelcontextprotocol.json.McpJsonMapper json;
50+
private boolean generateOutputSchema = false;
5051
@Override
5152
public void capabilities(io.modelcontextprotocol.spec.McpSchema.ServerCapabilities.Builder capabilities) {
5253
capabilities.tools(true);
@@ -55,6 +56,12 @@ public void capabilities(io.modelcontextprotocol.spec.McpSchema.ServerCapabiliti
5556
capabilities.completions();
5657
}
5758
59+
@Override
60+
public io.jooby.mcp.McpService generateOutputSchema(boolean generateOutputSchema) {
61+
this.generateOutputSchema = generateOutputSchema;
62+
return this;
63+
}
64+
5865
@Override
5966
public String serverKey() {
6067
return "example-server";

0 commit comments

Comments
 (0)