Skip to content

Commit 9cbdb09

Browse files
committed
This commit introduces the McpInvoker interface to wrap all generated
MCP tool, prompt, resource, and completion calls. It centralizes execution logic, exception handling, and protocol error mapping, while providing developers a clean hook to inject custom telemetry, tracing, or MDC context propagation. Details: * Added `McpInvoker` interface and `DefaultMcpInvoker` implementation. * Integrated Jooby's `Router.errorCode()` to seamlessly map standard framework exceptions (e.g., 400 Bad Request) to standard MCP JSON-RPC errors (e.g., -32602 Invalid Params). * Implemented LLM "self-healing" for tools: Unhandled business exceptions are now caught and returned as a `CallToolResult` with `isError=true`. This prevents protocol aborts and feeds the error text directly back into the LLM context so it can self-correct. * Updated the APT generator (`McpRouter`) to dynamically resolve the `McpInvoker` from the Jooby application registry using local variables, ensuring the generated router remains completely stateless. * Wrapped all routing lambdas in `invoker.invoke(operationId, action)`, passing contextual operation IDs (e.g., `tools/add_numbers` or `resources/calculator://history/{user}`).
1 parent 40a6e9b commit 9cbdb09

File tree

8 files changed

+298
-74
lines changed

8 files changed

+298
-74
lines changed

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

Lines changed: 105 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -240,9 +240,11 @@ public String toSourceCode(boolean kt) throws IOException {
240240
buffer.append(
241241
statement(
242242
indent(4),
243-
"override fun completions():"
243+
"override fun completions(app: io.jooby.Jooby):"
244244
+ " List<io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification>"
245245
+ " {"));
246+
buffer.append(
247+
statement(indent(6), "val invoker = app.require(io.jooby.mcp.McpInvoker::class.java)"));
246248
buffer.append(
247249
statement(
248250
indent(6),
@@ -255,7 +257,12 @@ public String toSourceCode(boolean kt) throws IOException {
255257
indent(4),
256258
"public"
257259
+ " java.util.List<io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification>"
258-
+ " completions() {"));
260+
+ " completions(io.jooby.Jooby app) {"));
261+
buffer.append(
262+
statement(
263+
indent(6),
264+
"var invoker = app.require(io.jooby.mcp.McpInvoker.class)",
265+
semicolon(kt)));
259266
buffer.append(
260267
statement(
261268
indent(6),
@@ -264,7 +271,6 @@ public String toSourceCode(boolean kt) throws IOException {
264271
semicolon(kt)));
265272
}
266273

267-
// Loop over ALL possible refs, not just the ones with explicit handlers
268274
for (var ref : allCompletionRefs) {
269275
var isResource = ref.contains("://");
270276
var refObj =
@@ -275,17 +281,20 @@ public String toSourceCode(boolean kt) throws IOException {
275281
String lambda;
276282
if (completionGroups.containsKey(ref)) {
277283
var handlerName = findTargetMethodName(ref) + "CompletionHandler";
284+
var operationId = "completions/" + ref;
278285
lambda =
279286
kt
280-
? "{ exchange, req -> this."
287+
? "{ exchange, req -> invoker.invoke("
288+
+ string(operationId)
289+
+ ") { this."
281290
+ handlerName
282-
+ "(exchange, exchange.transportContext(), req) }"
283-
: "(exchange, req) -> this."
291+
+ "(exchange, exchange.transportContext(), req) } }"
292+
: "(exchange, req) -> invoker.invoke("
293+
+ string(operationId)
294+
+ ", () -> this."
284295
+ handlerName
285-
+ "(exchange, exchange.transportContext(), req)";
286-
296+
+ "(exchange, exchange.transportContext(), req))";
287297
} else {
288-
// Fallback: Return an empty completion result safely
289298
lambda =
290299
kt
291300
? "{ _, _ -> io.jooby.mcp.McpResult(this.json).toCompleteResult(emptyList<Any>()) }"
@@ -328,9 +337,11 @@ public String toSourceCode(boolean kt) throws IOException {
328337
buffer.append(
329338
statement(
330339
indent(4),
331-
"override fun statelessCompletions():"
340+
"override fun statelessCompletions(app: io.jooby.Jooby):"
332341
+ " List<io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncCompletionSpecification>"
333342
+ " {"));
343+
buffer.append(
344+
statement(indent(6), "val invoker = app.require(io.jooby.mcp.McpInvoker::class.java)"));
334345
buffer.append(
335346
statement(
336347
indent(6),
@@ -343,7 +354,12 @@ public String toSourceCode(boolean kt) throws IOException {
343354
indent(4),
344355
"public"
345356
+ " java.util.List<io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncCompletionSpecification>"
346-
+ " statelessCompletions() {"));
357+
+ " statelessCompletions(io.jooby.Jooby app) {"));
358+
buffer.append(
359+
statement(
360+
indent(6),
361+
"var invoker = app.require(io.jooby.mcp.McpInvoker.class)",
362+
semicolon(kt)));
347363
buffer.append(
348364
statement(
349365
indent(6),
@@ -352,7 +368,6 @@ public String toSourceCode(boolean kt) throws IOException {
352368
semicolon(kt)));
353369
}
354370

355-
// Loop over ALL possible refs
356371
for (var ref : allCompletionRefs) {
357372
var isResource = ref.contains("://");
358373
var refObj =
@@ -363,12 +378,20 @@ public String toSourceCode(boolean kt) throws IOException {
363378
String lambda;
364379
if (completionGroups.containsKey(ref)) {
365380
var handlerName = findTargetMethodName(ref) + "CompletionHandler";
381+
var operationId = "completions/" + ref;
366382
lambda =
367383
kt
368-
? "{ ctx, req -> this." + handlerName + "(null, ctx, req) }"
369-
: "(ctx, req) -> this." + handlerName + "(null, ctx, req)";
384+
? "{ ctx, req -> invoker.invoke("
385+
+ string(operationId)
386+
+ ") { this."
387+
+ handlerName
388+
+ "(null, ctx, req) } }"
389+
: "(ctx, req) -> invoker.invoke("
390+
+ string(operationId)
391+
+ ", () -> this."
392+
+ handlerName
393+
+ "(null, ctx, req))";
370394
} else {
371-
// Fallback: Return an empty completion result safely
372395
lambda =
373396
kt
374397
? "{ _, _ -> io.jooby.mcp.McpResult(this.json).toCompleteResult(emptyList<Any>()) }"
@@ -424,6 +447,8 @@ public String toSourceCode(boolean kt) throws IOException {
424447
indent(6),
425448
"this.json ="
426449
+ " app.require(io.modelcontextprotocol.json.McpJsonMapper::class.java)"));
450+
buffer.append(
451+
statement(indent(6), "val invoker = app.require(io.jooby.mcp.McpInvoker::class.java)"));
427452

428453
if (!tools.isEmpty()) {
429454
buffer.append(
@@ -445,6 +470,11 @@ public String toSourceCode(boolean kt) throws IOException {
445470
indent(6),
446471
"this.json =" + " app.require(io.modelcontextprotocol.json.McpJsonMapper.class)",
447472
semicolon(kt)));
473+
buffer.append(
474+
statement(
475+
indent(6),
476+
"var invoker = app.require(io.jooby.mcp.McpInvoker.class)",
477+
semicolon(kt)));
448478

449479
if (!tools.isEmpty()) {
450480
buffer.append(
@@ -464,22 +494,72 @@ public String toSourceCode(boolean kt) throws IOException {
464494
for (var route : getRoutes()) {
465495
var methodName = route.getMethodName();
466496

497+
String mcpType = "";
498+
String mcpName = "";
499+
if (route.isMcpTool()) {
500+
mcpType = "tools";
501+
var ann =
502+
AnnotationSupport.findAnnotationByName(
503+
route.getMethod(), "io.jooby.annotation.mcp.McpTool");
504+
mcpName =
505+
ann != null
506+
? AnnotationSupport.findAnnotationValue(ann, "name"::equals).stream()
507+
.findFirst()
508+
.orElse("")
509+
: "";
510+
} else if (route.isMcpPrompt()) {
511+
mcpType = "prompts";
512+
var ann =
513+
AnnotationSupport.findAnnotationByName(
514+
route.getMethod(), "io.jooby.annotation.mcp.McpPrompt");
515+
mcpName =
516+
ann != null
517+
? AnnotationSupport.findAnnotationValue(ann, "name"::equals).stream()
518+
.findFirst()
519+
.orElse("")
520+
: "";
521+
} else if (route.isMcpResource() || route.isMcpResourceTemplate()) {
522+
mcpType = "resources";
523+
var ann =
524+
AnnotationSupport.findAnnotationByName(
525+
route.getMethod(), "io.jooby.annotation.mcp.McpResource");
526+
mcpName =
527+
ann != null
528+
? AnnotationSupport.findAnnotationValue(ann, "uri"::equals).stream()
529+
.findFirst()
530+
.orElse("")
531+
: "";
532+
}
533+
if (mcpName == null || mcpName.isEmpty()) mcpName = methodName;
534+
String operationId = mcpType + "/" + mcpName;
535+
467536
// --- Lambda Router Definition ---
468537
String lambda =
469538
kt
470539
? (isStateless
471-
? "{ ctx, req -> this." + methodName + "(null, ctx, req) }"
472-
: "{ exchange, req -> this."
540+
? "{ ctx, req -> invoker.invoke("
541+
+ string(operationId)
542+
+ ") { this."
543+
+ methodName
544+
+ "(null, ctx, req) } }"
545+
: "{ exchange, req -> invoker.invoke("
546+
+ string(operationId)
547+
+ ") { this."
473548
+ methodName
474-
+ "(exchange, exchange.transportContext(), req) }")
549+
+ "(exchange, exchange.transportContext(), req) } }")
475550
: (isStateless
476-
? "(ctx, req) -> this." + methodName + "(null, ctx, req)"
477-
: "(exchange, req) -> this."
551+
? "(ctx, req) -> invoker.invoke("
552+
+ string(operationId)
553+
+ ", () -> this."
478554
+ methodName
479-
+ "(exchange, exchange.transportContext(), req)");
555+
+ "(null, ctx, req))"
556+
: "(exchange, req) -> invoker.invoke("
557+
+ string(operationId)
558+
+ ", () -> this."
559+
+ methodName
560+
+ "(exchange, exchange.transportContext(), req))");
480561

481562
if (route.isMcpTool()) {
482-
// Removed "mapper" from defArgs
483563
var defArgs = "schemaGenerator";
484564
if (kt) {
485565
buffer.append(
@@ -593,19 +673,11 @@ public String toSourceCode(boolean kt) throws IOException {
593673
"private fun ",
594674
handlerName,
595675
"(exchange: io.modelcontextprotocol.server.McpSyncServerExchange?,"
596-
+ " transportContext:"
597-
+ " io.modelcontextprotocol.common.McpTransportContext,"
598-
+ " req:" // Removed '?'
676+
+ " transportContext: io.modelcontextprotocol.common.McpTransportContext, req:"
599677
+ " io.modelcontextprotocol.spec.McpSchema.CompleteRequest):"
600678
+ " io.modelcontextprotocol.spec.McpSchema.CompleteResult {"));
601-
602-
// Direct extraction, no fallback needed
603679
buffer.append(
604-
statement(
605-
indent(6),
606-
"val ctx ="
607-
+ " transportContext.get<io.jooby.Context>(io.jooby.Context::class.java.name)"));
608-
680+
statement(indent(6), "val ctx = transportContext.get(\"CTX\") as io.jooby.Context"));
609681
buffer.append(statement(indent(6), "val c = this.factory.apply(ctx)"));
610682
buffer.append(statement(indent(6), "val targetArg = req.argument()?.name() ?: \"\""));
611683
buffer.append(statement(indent(6), "val typedValue = req.argument()?.value() ?: \"\""));
@@ -617,17 +689,13 @@ public String toSourceCode(boolean kt) throws IOException {
617689
"private io.modelcontextprotocol.spec.McpSchema.CompleteResult ",
618690
handlerName,
619691
"(io.modelcontextprotocol.server.McpSyncServerExchange exchange,"
620-
+ " io.modelcontextprotocol.common.McpTransportContext"
621-
+ " transportContext," // Guaranteed non-null
692+
+ " io.modelcontextprotocol.common.McpTransportContext transportContext,"
622693
+ " io.modelcontextprotocol.spec.McpSchema.CompleteRequest req) {"));
623-
624-
// Direct extraction, no ternary operator
625694
buffer.append(
626695
statement(
627696
indent(6),
628697
"var ctx = (io.jooby.Context) transportContext.get(\"CTX\")",
629698
semicolon(kt)));
630-
631699
buffer.append(statement(indent(6), "var c = this.factory.apply(ctx)", semicolon(kt)));
632700
buffer.append(
633701
statement(

0 commit comments

Comments
 (0)