Skip to content

Commit 40a6e9b

Browse files
jknackjknack
committed
refactor(mcp): simplify transport context extraction in code generator
Since the routing lambdas now directly inject the transport context for both stateful and stateless flows, the `transportContext` parameter is guaranteed to be non-null in the generated handler methods. This commit cleans up the APT generator (`McpRoute`, `McpRouter`) by: * Removing redundant null-checks and ternary operators associated with the transport context. * Eliminating Kotlin safe-calls (`?.`) in method signatures and variable extractions. * Streamlining Jooby `Context` extraction to use direct, safe casts (e.g., `(io.jooby.Context) transportContext.get("CTX")`). * Significantly reducing the boilerplate and branching in all generated tool, prompt, resource, and completion handlers.
1 parent d33b90b commit 40a6e9b

File tree

5 files changed

+63
-63
lines changed

5 files changed

+63
-63
lines changed

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

Lines changed: 9 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -708,18 +708,15 @@ public List<String> generateMcpHandlerMethod(boolean kt) {
708708
"private fun ",
709709
handlerName,
710710
"(exchange: io.modelcontextprotocol.server.McpSyncServerExchange?, transportContext:"
711-
+ " io.modelcontextprotocol.common.McpTransportContext?, req:"
711+
+ " io.modelcontextprotocol.common.McpTransportContext, req:" // Removed '?'
712712
+ " io.modelcontextprotocol.spec.McpSchema.",
713713
reqType,
714714
"): io.modelcontextprotocol.spec.McpSchema.",
715715
resType,
716716
" {"));
717+
717718
buffer.add(
718-
statement(
719-
indent(6),
720-
"val ctx ="
721-
+ " exchange?.transportContext()?.get<io.jooby.Context>(io.jooby.Context::class.java.name)"
722-
+ " ?: transportContext?.get<io.jooby.Context>(io.jooby.Context::class.java.name)"));
719+
statement(indent(6), "val ctx = transportContext.get(\"CTX\") as? io.jooby.Context"));
723720
} else {
724721
buffer.add(
725722
statement(
@@ -729,16 +726,16 @@ public List<String> generateMcpHandlerMethod(boolean kt) {
729726
" ",
730727
handlerName,
731728
"(io.modelcontextprotocol.server.McpSyncServerExchange exchange,"
732-
+ " io.modelcontextprotocol.common.McpTransportContext transportContext,"
729+
+ " io.modelcontextprotocol.common.McpTransportContext"
730+
+ " transportContext," // Guaranteed non-null
733731
+ " io.modelcontextprotocol.spec.McpSchema.",
734732
reqType,
735733
" req) {"));
734+
736735
buffer.add(
737736
statement(
738737
indent(6),
739-
"var ctx = exchange != null ? (io.jooby.Context)"
740-
+ " exchange.transportContext().get(\"CTX\") : (transportContext != null ?"
741-
+ " (io.jooby.Context) transportContext.get(\"CTX\") : null)",
738+
"var ctx = (io.jooby.Context) transportContext.get(\"CTX\")",
742739
semicolon(kt)));
743740
}
744741

@@ -809,26 +806,8 @@ public List<String> generateMcpHandlerMethod(boolean kt) {
809806
javaParamNames.add(javaName);
810807

811808
if (type.equals("io.jooby.Context")
812-
|| type.equals("io.modelcontextprotocol.server.McpSyncServerExchange")) {
813-
continue;
814-
}
815-
if (type.equals("io.modelcontextprotocol.common.McpTransportContext")) {
816-
if (kt) {
817-
buffer.add(
818-
statement(
819-
indent(6),
820-
"val ",
821-
javaName,
822-
" = exchange?.transportContext() ?: transportContext"));
823-
} else {
824-
buffer.add(
825-
statement(
826-
indent(6),
827-
"var ",
828-
javaName,
829-
" = exchange != null ? exchange.transportContext() : transportContext",
830-
semicolon(kt)));
831-
}
809+
|| type.equals("io.modelcontextprotocol.server.McpSyncServerExchange")
810+
|| type.equals("io.modelcontextprotocol.common.McpTransportContext")) {
832811
continue;
833812
} else if (type.equals("io.modelcontextprotocol.spec.McpSchema." + reqType)) {
834813
buffer.add(statement(indent(6), kt ? "val " : "var ", javaName, " = req", semicolon(kt)));

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

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -277,8 +277,13 @@ public String toSourceCode(boolean kt) throws IOException {
277277
var handlerName = findTargetMethodName(ref) + "CompletionHandler";
278278
lambda =
279279
kt
280-
? "{ exchange, req -> this." + handlerName + "(exchange, null, req) }"
281-
: "(exchange, req) -> this." + handlerName + "(exchange, null, req)";
280+
? "{ exchange, req -> this."
281+
+ handlerName
282+
+ "(exchange, exchange.transportContext(), req) }"
283+
: "(exchange, req) -> this."
284+
+ handlerName
285+
+ "(exchange, exchange.transportContext(), req)";
286+
282287
} else {
283288
// Fallback: Return an empty completion result safely
284289
lambda =
@@ -464,10 +469,14 @@ public String toSourceCode(boolean kt) throws IOException {
464469
kt
465470
? (isStateless
466471
? "{ ctx, req -> this." + methodName + "(null, ctx, req) }"
467-
: "{ exchange, req -> this." + methodName + "(exchange, null, req) }")
472+
: "{ exchange, req -> this."
473+
+ methodName
474+
+ "(exchange, exchange.transportContext(), req) }")
468475
: (isStateless
469476
? "(ctx, req) -> this." + methodName + "(null, ctx, req)"
470-
: "(exchange, req) -> this." + methodName + "(exchange, null, req)");
477+
: "(exchange, req) -> this."
478+
+ methodName
479+
+ "(exchange, exchange.transportContext(), req)");
471480

472481
if (route.isMcpTool()) {
473482
// Removed "mapper" from defArgs
@@ -584,15 +593,19 @@ public String toSourceCode(boolean kt) throws IOException {
584593
"private fun ",
585594
handlerName,
586595
"(exchange: io.modelcontextprotocol.server.McpSyncServerExchange?,"
587-
+ " transportContext: io.modelcontextprotocol.common.McpTransportContext?, req:"
596+
+ " transportContext:"
597+
+ " io.modelcontextprotocol.common.McpTransportContext,"
598+
+ " req:" // Removed '?'
588599
+ " io.modelcontextprotocol.spec.McpSchema.CompleteRequest):"
589600
+ " io.modelcontextprotocol.spec.McpSchema.CompleteResult {"));
601+
602+
// Direct extraction, no fallback needed
590603
buffer.append(
591604
statement(
592605
indent(6),
593606
"val ctx ="
594-
+ " exchange?.transportContext()?.get<io.jooby.Context>(io.jooby.Context::class.java.name)"
595-
+ " ?: transportContext?.get<io.jooby.Context>(io.jooby.Context::class.java.name)"));
607+
+ " transportContext.get<io.jooby.Context>(io.jooby.Context::class.java.name)"));
608+
596609
buffer.append(statement(indent(6), "val c = this.factory.apply(ctx)"));
597610
buffer.append(statement(indent(6), "val targetArg = req.argument()?.name() ?: \"\""));
598611
buffer.append(statement(indent(6), "val typedValue = req.argument()?.value() ?: \"\""));
@@ -604,15 +617,17 @@ public String toSourceCode(boolean kt) throws IOException {
604617
"private io.modelcontextprotocol.spec.McpSchema.CompleteResult ",
605618
handlerName,
606619
"(io.modelcontextprotocol.server.McpSyncServerExchange exchange,"
607-
+ " io.modelcontextprotocol.common.McpTransportContext transportContext,"
620+
+ " io.modelcontextprotocol.common.McpTransportContext"
621+
+ " transportContext," // Guaranteed non-null
608622
+ " io.modelcontextprotocol.spec.McpSchema.CompleteRequest req) {"));
623+
624+
// Direct extraction, no ternary operator
609625
buffer.append(
610626
statement(
611627
indent(6),
612-
"var ctx = exchange != null ? (io.jooby.Context)"
613-
+ " exchange.transportContext().get(\"CTX\") : (transportContext != null ?"
614-
+ " (io.jooby.Context) transportContext.get(\"CTX\") : null)",
628+
"var ctx = (io.jooby.Context) transportContext.get(\"CTX\")",
615629
semicolon(kt)));
630+
616631
buffer.append(statement(indent(6), "var c = this.factory.apply(ctx)", semicolon(kt)));
617632
buffer.append(
618633
statement(
@@ -638,11 +653,7 @@ public String toSourceCode(boolean kt) throws IOException {
638653
} else if (type.equals("io.modelcontextprotocol.server.McpSyncServerExchange")) {
639654
invokeArgs.add("exchange");
640655
} else if (type.equals("io.modelcontextprotocol.common.McpTransportContext")) {
641-
if (kt) {
642-
invokeArgs.add("exchange?.transportContext() ?: transportContext");
643-
} else {
644-
invokeArgs.add("exchange != null ? exchange.transportContext() : transportContext");
645-
}
656+
invokeArgs.add("transportContext");
646657
} else {
647658
targetArgName = param.getMcpName();
648659
invokeArgs.add("typedValue");

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

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ public String serverKey() {
6363
@Override
6464
public java.util.List<io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification> completions() {
6565
var completions = new java.util.ArrayList<io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification>();
66-
completions.add(new io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.PromptReference("review_code"), (exchange, req) -> this.reviewCodeCompletionHandler(exchange, null, req)));
67-
completions.add(new io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.ResourceReference("file:///users/{id}/{name}/profile"), (exchange, req) -> this.getUserProfileCompletionHandler(exchange, null, req)));
66+
completions.add(new io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.PromptReference("review_code"), (exchange, req) -> this.reviewCodeCompletionHandler(exchange, exchange.transportContext(), req)));
67+
completions.add(new io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.ResourceReference("file:///users/{id}/{name}/profile"), (exchange, req) -> this.getUserProfileCompletionHandler(exchange, exchange.transportContext(), req)));
6868
return completions;
6969
}
7070
@@ -81,10 +81,10 @@ public void install(io.jooby.Jooby app, io.modelcontextprotocol.server.McpSyncSe
8181
this.json = app.require(io.modelcontextprotocol.json.McpJsonMapper.class);
8282
var schemaGenerator = app.require(com.github.victools.jsonschema.generator.SchemaGenerator.class);
8383
84-
server.addTool(new io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification(addToolSpec(schemaGenerator), (exchange, req) -> this.add(exchange, null, req)));
85-
server.addPrompt(new io.modelcontextprotocol.server.McpServerFeatures.SyncPromptSpecification(reviewCodePromptSpec(), (exchange, req) -> this.reviewCode(exchange, null, req)));
86-
server.addResource(new io.modelcontextprotocol.server.McpServerFeatures.SyncResourceSpecification(getLogsResourceSpec(), (exchange, req) -> this.getLogs(exchange, null, req)));
87-
server.addResourceTemplate(new io.modelcontextprotocol.server.McpServerFeatures.SyncResourceTemplateSpecification(getUserProfileResourceTemplateSpec(), (exchange, req) -> this.getUserProfile(exchange, null, req)));
84+
server.addTool(new io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification(addToolSpec(schemaGenerator), (exchange, req) -> this.add(exchange, exchange.transportContext(), req)));
85+
server.addPrompt(new io.modelcontextprotocol.server.McpServerFeatures.SyncPromptSpecification(reviewCodePromptSpec(), (exchange, req) -> this.reviewCode(exchange, exchange.transportContext(), req)));
86+
server.addResource(new io.modelcontextprotocol.server.McpServerFeatures.SyncResourceSpecification(getLogsResourceSpec(), (exchange, req) -> this.getLogs(exchange, exchange.transportContext(), req)));
87+
server.addResourceTemplate(new io.modelcontextprotocol.server.McpServerFeatures.SyncResourceTemplateSpecification(getUserProfileResourceTemplateSpec(), (exchange, req) -> this.getUserProfile(exchange, exchange.transportContext(), req)));
8888
}
8989
9090
@Override
@@ -118,7 +118,7 @@ private io.modelcontextprotocol.spec.McpSchema.Tool addToolSpec(com.github.victo
118118
}
119119
120120
private io.modelcontextprotocol.spec.McpSchema.CallToolResult add(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.CallToolRequest req) {
121-
var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
121+
var ctx = (io.jooby.Context) transportContext.get("CTX");
122122
var args = req.arguments() != null ? req.arguments() : java.util.Collections.<String, Object>emptyMap();
123123
var c = this.factory.apply(ctx);
124124
var raw_a = args.get("a");
@@ -139,7 +139,7 @@ private io.modelcontextprotocol.spec.McpSchema.Prompt reviewCodePromptSpec() {
139139
}
140140
141141
private io.modelcontextprotocol.spec.McpSchema.GetPromptResult reviewCode(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.GetPromptRequest req) {
142-
var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
142+
var ctx = (io.jooby.Context) transportContext.get("CTX");
143143
var args = req.arguments() != null ? req.arguments() : java.util.Collections.<String, Object>emptyMap();
144144
var c = this.factory.apply(ctx);
145145
var raw_language = args.get("language");
@@ -157,7 +157,7 @@ private io.modelcontextprotocol.spec.McpSchema.Resource getLogsResourceSpec() {
157157
}
158158
159159
private io.modelcontextprotocol.spec.McpSchema.ReadResourceResult getLogs(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.ReadResourceRequest req) {
160-
var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
160+
var ctx = (io.jooby.Context) transportContext.get("CTX");
161161
var args = java.util.Collections.<String, Object>emptyMap();
162162
var c = this.factory.apply(ctx);
163163
var result = c.getLogs();
@@ -169,7 +169,7 @@ private io.modelcontextprotocol.spec.McpSchema.ResourceTemplate getUserProfileRe
169169
}
170170
171171
private io.modelcontextprotocol.spec.McpSchema.ReadResourceResult getUserProfile(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.ReadResourceRequest req) {
172-
var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
172+
var ctx = (io.jooby.Context) transportContext.get("CTX");
173173
var uri = req.uri();
174174
var manager = new io.modelcontextprotocol.util.DefaultMcpUriTemplateManager("file:///users/{id}/{name}/profile");
175175
var args = new java.util.HashMap<String, Object>();
@@ -184,7 +184,7 @@ private io.modelcontextprotocol.spec.McpSchema.ReadResourceResult getUserProfile
184184
}
185185
186186
private io.modelcontextprotocol.spec.McpSchema.CompleteResult getUserProfileCompletionHandler(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.CompleteRequest req) {
187-
var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
187+
var ctx = (io.jooby.Context) transportContext.get("CTX");
188188
var c = this.factory.apply(ctx);
189189
var targetArg = req.argument() != null ? req.argument().name() : "";
190190
var typedValue = req.argument() != null ? req.argument().value() : "";
@@ -202,7 +202,7 @@ private io.modelcontextprotocol.spec.McpSchema.CompleteResult getUserProfileComp
202202
}
203203
204204
private io.modelcontextprotocol.spec.McpSchema.CompleteResult reviewCodeCompletionHandler(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.CompleteRequest req) {
205-
var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
205+
var ctx = (io.jooby.Context) transportContext.get("CTX");
206206
var c = this.factory.apply(ctx);
207207
var targetArg = req.argument() != null ? req.argument().name() : "";
208208
var typedValue = req.argument() != null ? req.argument().value() : "";

‎modules/jooby-mcp/src/main/java/io/jooby/mcp/McpHandler.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
/*
2+
* Jooby https://jooby.io
3+
* Apache License Version 2.0 https://jooby.io/LICENSE.txt
4+
* Copyright 2014 Edgar Espina
5+
*/
6+
package io.jooby.mcp;
7+
8+
import io.jooby.SneakyThrows;
9+
10+
public class McpHandler {
11+
12+
public static <R> R invoke(SneakyThrows.Supplier<R> action) {
13+
return action.get();
14+
}
15+
}

‎tests/src/test/java/io/jooby/i3830/CalculatorTools.java‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,6 @@
1717
/** A collection of tools, prompts, and resources exposed to the LLM via MCP. */
1818
public class CalculatorTools {
1919

20-
// --- TOOLS ---
21-
2220
/**
2321
* Adds two integers together and returns the result.
2422
*
@@ -31,15 +29,13 @@ public int add(int a, int b) {
3129
return a + b;
3230
}
3331

34-
// --- PROMPTS ---
3532
@McpPrompt(name = "math_tutor", description = "A prompt to initiate a math tutoring session")
3633
public String mathTutor(String topic) {
3734
return "You are a helpful math tutor. Please explain the concept of "
3835
+ topic
3936
+ " step by step.";
4037
}
4138

42-
// --- RESOURCES ---
4339
@McpResource(
4440
uri = "calculator://manual/usage",
4541
name = "Calculator Manual",
@@ -56,7 +52,6 @@ public String history(String user) {
5652
return "History for " + user + ":\n5 + 10 = 15\n2 * 4 = 8";
5753
}
5854

59-
// --- COMPLETIONS ---
6055
@McpCompletion(ref = "calculator://history/{user}")
6156
public List<String> historyCompletions(String user) {
6257
// In a real app, this would query a database for active usernames matching the input

0 commit comments

Comments
 (0)