Skip to content

Commit c1cde20

Browse files
tzolovtzolov
authored
refactor: improve MCP server resilience and error handling (#579)
- Replace McpError with appropriate standard exceptions (IllegalArgumentException, IllegalStateException) for validation errors - Change tool/prompt registration to replace existing items instead of rejecting duplicates with warningsi - Change addTool behavior to replace existing tools instead of throwing errors - Change addPrompt() to allow replacing existing prompts with warning instead of throwing error - Make removal operations idempotent - log warnings instead of throwing errors for non-existent items - Change removePrompt() to log warning instead of throwing error for non-existent prompts - Change removeTool() to gracefully handle non-existent tools with warnings instead of errors - Add listTools() and listPrompts() methods to all server variants (async, sync, stateless) - Add listTools() method to all server classes for tool enumeration - Add listPrompts() method to all server implementations for retrieving registered prompts - Improve error construction using McpError.builder() pattern for protocol-specific errors - Update error handling in session classes to properly propagate McpError JSON-RPC errors - Use proper error codes (INVALID_PARAMS) for prompt not found scenarios - Update tests to reflect new lenient behavior for duplicate registrations and removals - Add aggregateExceptionMessages() utility method utility to aggregate the exception chains - Update McpServerSession and McpStreamableServerSession to use the aggregated errors in the json-rpc-error data section This change makes the MCP server APIs more resilient and user-friendly by using appropriate exception types, supporting item replacement, and providing listing capabilities while maintaining backward compatibility. Signed-off-by: Christian Tzolov <[email protected]>
1 parent ff98e29 commit c1cde20

File tree

13 files changed

+305
-196
lines changed

13 files changed

+305
-196
lines changed

‎mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java‎

Lines changed: 57 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -322,25 +322,24 @@ private McpNotificationHandler asyncRootsListChangedNotificationHandler(
322322
*/
323323
public Mono<Void> addTool(McpServerFeatures.AsyncToolSpecification toolSpecification) {
324324
if (toolSpecification == null) {
325-
return Mono.error(new McpError("Tool specification must not be null"));
325+
return Mono.error(new IllegalArgumentException("Tool specification must not be null"));
326326
}
327327
if (toolSpecification.tool() == null) {
328-
return Mono.error(new McpError("Tool must not be null"));
328+
return Mono.error(new IllegalArgumentException("Tool must not be null"));
329329
}
330330
if (toolSpecification.call() == null && toolSpecification.callHandler() == null) {
331-
return Mono.error(new McpError("Tool call handler must not be null"));
331+
return Mono.error(new IllegalArgumentException("Tool call handler must not be null"));
332332
}
333333
if (this.serverCapabilities.tools() == null) {
334-
return Mono.error(new McpError("Server must be configured with tool capabilities"));
334+
return Mono.error(new IllegalStateException("Server must be configured with tool capabilities"));
335335
}
336336

337337
var wrappedToolSpecification = withStructuredOutputHandling(this.jsonSchemaValidator, toolSpecification);
338338

339339
return Mono.defer(() -> {
340-
// Check for duplicate tool names
341-
if (this.tools.stream().anyMatch(th -> th.tool().name().equals(wrappedToolSpecification.tool().name()))) {
342-
return Mono.error(
343-
new McpError("Tool with name '" + wrappedToolSpecification.tool().name() + "' already exists"));
340+
// Remove tools with duplicate tool names first
341+
if (this.tools.removeIf(th -> th.tool().name().equals(wrappedToolSpecification.tool().name()))) {
342+
logger.warn("Replace existing Tool with name '{}'", wrappedToolSpecification.tool().name());
344343
}
345344

346345
this.tools.add(wrappedToolSpecification);
@@ -464,30 +463,40 @@ private static McpServerFeatures.AsyncToolSpecification withStructuredOutputHand
464463
.build();
465464
}
466465

466+
/**
467+
* List all registered tools.
468+
* @return A Flux stream of all registered tools
469+
*/
470+
public Flux<Tool> listTools() {
471+
return Flux.fromIterable(this.tools).map(McpServerFeatures.AsyncToolSpecification::tool);
472+
}
473+
467474
/**
468475
* Remove a tool handler at runtime.
469476
* @param toolName The name of the tool handler to remove
470477
* @return Mono that completes when clients have been notified of the change
471478
*/
472479
public Mono<Void> removeTool(String toolName) {
473480
if (toolName == null) {
474-
return Mono.error(new McpError("Tool name must not be null"));
481+
return Mono.error(new IllegalArgumentException("Tool name must not be null"));
475482
}
476483
if (this.serverCapabilities.tools() == null) {
477-
return Mono.error(new McpError("Server must be configured with tool capabilities"));
484+
return Mono.error(new IllegalStateException("Server must be configured with tool capabilities"));
478485
}
479486

480487
return Mono.defer(() -> {
481-
boolean removed = this.tools
482-
.removeIf(toolSpecification -> toolSpecification.tool().name().equals(toolName));
483-
if (removed) {
488+
if (this.tools.removeIf(toolSpecification -> toolSpecification.tool().name().equals(toolName))) {
489+
484490
logger.debug("Removed tool handler: {}", toolName);
485491
if (this.serverCapabilities.tools().listChanged()) {
486492
return notifyToolsListChanged();
487493
}
488-
return Mono.empty();
489494
}
490-
return Mono.error(new McpError("Tool with name '" + toolName + "' not found"));
495+
else {
496+
logger.warn("Ignore as a Tool with name '{}' not found", toolName);
497+
}
498+
499+
return Mono.empty();
491500
});
492501
}
493502

@@ -518,8 +527,10 @@ private McpRequestHandler<CallToolResult> toolsCallRequestHandler() {
518527
.findAny();
519528

520529
if (toolSpecification.isEmpty()) {
521-
return Mono.error(new McpError(new JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INVALID_PARAMS,
522-
"Unknown tool: invalid_tool_name", "Tool not found: " + callToolRequest.name())));
530+
return Mono.error(McpError.builder(McpSchema.ErrorCodes.INVALID_PARAMS)
531+
.message("Unknown tool: invalid_tool_name")
532+
.data("Tool not found: " + callToolRequest.name())
533+
.build());
523534
}
524535

525536
return toolSpecification.get().callHandler().apply(exchange, callToolRequest);
@@ -747,58 +758,63 @@ private Optional<McpServerFeatures.AsyncResourceTemplateSpecification> findResou
747758
*/
748759
public Mono<Void> addPrompt(McpServerFeatures.AsyncPromptSpecification promptSpecification) {
749760
if (promptSpecification == null) {
750-
return Mono.error(new McpError("Prompt specification must not be null"));
761+
return Mono.error(new IllegalArgumentException("Prompt specification must not be null"));
751762
}
752763
if (this.serverCapabilities.prompts() == null) {
753-
return Mono.error(new McpError("Server must be configured with prompt capabilities"));
764+
return Mono.error(new IllegalStateException("Server must be configured with prompt capabilities"));
754765
}
755766

756767
return Mono.defer(() -> {
757-
McpServerFeatures.AsyncPromptSpecification specification = this.prompts
758-
.putIfAbsent(promptSpecification.prompt().name(), promptSpecification);
759-
if (specification != null) {
760-
return Mono.error(
761-
new McpError("Prompt with name '" + promptSpecification.prompt().name() + "' already exists"));
768+
var previous = this.prompts.put(promptSpecification.prompt().name(), promptSpecification);
769+
if (previous != null) {
770+
logger.warn("Replace existing Prompt with name '{}'", promptSpecification.prompt().name());
771+
}
772+
else {
773+
logger.debug("Added prompt handler: {}", promptSpecification.prompt().name());
762774
}
763-
764-
logger.debug("Added prompt handler: {}", promptSpecification.prompt().name());
765-
766-
// Servers that declared the listChanged capability SHOULD send a
767-
// notification,
768-
// when the list of available prompts changes
769775
if (this.serverCapabilities.prompts().listChanged()) {
770-
return notifyPromptsListChanged();
776+
return this.notifyPromptsListChanged();
771777
}
778+
772779
return Mono.empty();
773780
});
774781
}
775782

783+
/**
784+
* List all registered prompts.
785+
* @return A Flux stream of all registered prompts
786+
*/
787+
public Flux<McpSchema.Prompt> listPrompts() {
788+
return Flux.fromIterable(this.prompts.values()).map(McpServerFeatures.AsyncPromptSpecification::prompt);
789+
}
790+
776791
/**
777792
* Remove a prompt handler at runtime.
778793
* @param promptName The name of the prompt handler to remove
779794
* @return Mono that completes when clients have been notified of the change
780795
*/
781796
public Mono<Void> removePrompt(String promptName) {
782797
if (promptName == null) {
783-
return Mono.error(new McpError("Prompt name must not be null"));
798+
return Mono.error(new IllegalArgumentException("Prompt name must not be null"));
784799
}
785800
if (this.serverCapabilities.prompts() == null) {
786-
return Mono.error(new McpError("Server must be configured with prompt capabilities"));
801+
return Mono.error(new IllegalStateException("Server must be configured with prompt capabilities"));
787802
}
788803

789804
return Mono.defer(() -> {
790805
McpServerFeatures.AsyncPromptSpecification removed = this.prompts.remove(promptName);
791806

792807
if (removed != null) {
793808
logger.debug("Removed prompt handler: {}", promptName);
794-
// Servers that declared the listChanged capability SHOULD send a
795-
// notification, when the list of available prompts changes
796809
if (this.serverCapabilities.prompts().listChanged()) {
797810
return this.notifyPromptsListChanged();
798811
}
799812
return Mono.empty();
800813
}
801-
return Mono.error(new McpError("Prompt with name '" + promptName + "' not found"));
814+
else {
815+
logger.warn("Ignore as a Prompt with name '{}' not found", promptName);
816+
}
817+
return Mono.empty();
802818
});
803819
}
804820

@@ -834,8 +850,12 @@ private McpRequestHandler<McpSchema.GetPromptResult> promptsGetRequestHandler()
834850

835851
// Implement prompt retrieval logic here
836852
McpServerFeatures.AsyncPromptSpecification specification = this.prompts.get(promptRequest.name());
853+
837854
if (specification == null) {
838-
return Mono.error(new McpError("Prompt not found: " + promptRequest.name()));
855+
return Mono.error(McpError.builder(ErrorCodes.INVALID_PARAMS)
856+
.message("Invalid prompt name")
857+
.data("Prompt not found: " + promptRequest.name())
858+
.build());
839859
}
840860

841861
return Mono.defer(() -> specification.promptHandler().apply(exchange, promptRequest));

‎mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java‎

Lines changed: 56 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -319,25 +319,24 @@ public Mono<CallToolResult> apply(McpTransportContext transportContext, McpSchem
319319
*/
320320
public Mono<Void> addTool(McpStatelessServerFeatures.AsyncToolSpecification toolSpecification) {
321321
if (toolSpecification == null) {
322-
return Mono.error(new McpError("Tool specification must not be null"));
322+
return Mono.error(new IllegalArgumentException("Tool specification must not be null"));
323323
}
324324
if (toolSpecification.tool() == null) {
325-
return Mono.error(new McpError("Tool must not be null"));
325+
return Mono.error(new IllegalArgumentException("Tool must not be null"));
326326
}
327327
if (toolSpecification.callHandler() == null) {
328-
return Mono.error(new McpError("Tool call handler must not be null"));
328+
return Mono.error(new IllegalArgumentException("Tool call handler must not be null"));
329329
}
330330
if (this.serverCapabilities.tools() == null) {
331-
return Mono.error(new McpError("Server must be configured with tool capabilities"));
331+
return Mono.error(new IllegalStateException("Server must be configured with tool capabilities"));
332332
}
333333

334334
var wrappedToolSpecification = withStructuredOutputHandling(this.jsonSchemaValidator, toolSpecification);
335335

336336
return Mono.defer(() -> {
337-
// Check for duplicate tool names
338-
if (this.tools.stream().anyMatch(th -> th.tool().name().equals(wrappedToolSpecification.tool().name()))) {
339-
return Mono.error(
340-
new McpError("Tool with name '" + wrappedToolSpecification.tool().name() + "' already exists"));
337+
// Remove tools with duplicate tool names first
338+
if (this.tools.removeIf(th -> th.tool().name().equals(wrappedToolSpecification.tool().name()))) {
339+
logger.warn("Replace existing Tool with name '{}'", wrappedToolSpecification.tool().name());
341340
}
342341

343342
this.tools.add(wrappedToolSpecification);
@@ -347,27 +346,37 @@ public Mono<Void> addTool(McpStatelessServerFeatures.AsyncToolSpecification tool
347346
});
348347
}
349348

349+
/**
350+
* List all registered tools.
351+
* @return A Flux stream of all registered tools
352+
*/
353+
public Flux<Tool> listTools() {
354+
return Flux.fromIterable(this.tools).map(McpStatelessServerFeatures.AsyncToolSpecification::tool);
355+
}
356+
350357
/**
351358
* Remove a tool handler at runtime.
352359
* @param toolName The name of the tool handler to remove
353360
* @return Mono that completes when clients have been notified of the change
354361
*/
355362
public Mono<Void> removeTool(String toolName) {
356363
if (toolName == null) {
357-
return Mono.error(new McpError("Tool name must not be null"));
364+
return Mono.error(new IllegalArgumentException("Tool name must not be null"));
358365
}
359366
if (this.serverCapabilities.tools() == null) {
360-
return Mono.error(new McpError("Server must be configured with tool capabilities"));
367+
return Mono.error(new IllegalStateException("Server must be configured with tool capabilities"));
361368
}
362369

363370
return Mono.defer(() -> {
364-
boolean removed = this.tools
365-
.removeIf(toolSpecification -> toolSpecification.tool().name().equals(toolName));
366-
if (removed) {
371+
if (this.tools.removeIf(toolSpecification -> toolSpecification.tool().name().equals(toolName))) {
372+
367373
logger.debug("Removed tool handler: {}", toolName);
368-
return Mono.empty();
369374
}
370-
return Mono.error(new McpError("Tool with name '" + toolName + "' not found"));
375+
else {
376+
logger.warn("Ignore as a Tool with name '{}' not found", toolName);
377+
}
378+
379+
return Mono.empty();
371380
});
372381
}
373382

@@ -391,8 +400,10 @@ private McpStatelessRequestHandler<CallToolResult> toolsCallRequestHandler() {
391400
.findAny();
392401

393402
if (toolSpecification.isEmpty()) {
394-
return Mono.error(new McpError(new JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INVALID_PARAMS,
395-
"Unknown tool: invalid_tool_name", "Tool not found: " + callToolRequest.name())));
403+
return Mono.error(McpError.builder(McpSchema.ErrorCodes.INVALID_PARAMS)
404+
.message("Unknown tool: invalid_tool_name")
405+
.data("Tool not found: " + callToolRequest.name())
406+
.build());
396407
}
397408

398409
return toolSpecification.get().callHandler().apply(ctx, callToolRequest);
@@ -593,37 +604,45 @@ private Optional<McpStatelessServerFeatures.AsyncResourceTemplateSpecification>
593604
*/
594605
public Mono<Void> addPrompt(McpStatelessServerFeatures.AsyncPromptSpecification promptSpecification) {
595606
if (promptSpecification == null) {
596-
return Mono.error(new McpError("Prompt specification must not be null"));
607+
return Mono.error(new IllegalArgumentException("Prompt specification must not be null"));
597608
}
598609
if (this.serverCapabilities.prompts() == null) {
599-
return Mono.error(new McpError("Server must be configured with prompt capabilities"));
610+
return Mono.error(new IllegalStateException("Server must be configured with prompt capabilities"));
600611
}
601612

602613
return Mono.defer(() -> {
603-
McpStatelessServerFeatures.AsyncPromptSpecification specification = this.prompts
604-
.putIfAbsent(promptSpecification.prompt().name(), promptSpecification);
605-
if (specification != null) {
606-
return Mono.error(
607-
new McpError("Prompt with name '" + promptSpecification.prompt().name() + "' already exists"));
614+
var previous = this.prompts.put(promptSpecification.prompt().name(), promptSpecification);
615+
if (previous != null) {
616+
logger.warn("Replace existing Prompt with name '{}'", promptSpecification.prompt().name());
617+
}
618+
else {
619+
logger.debug("Added prompt handler: {}", promptSpecification.prompt().name());
608620
}
609-
610-
logger.debug("Added prompt handler: {}", promptSpecification.prompt().name());
611621

612622
return Mono.empty();
613623
});
614624
}
615625

626+
/**
627+
* List all registered prompts.
628+
* @return A Flux stream of all registered prompts
629+
*/
630+
public Flux<McpSchema.Prompt> listPrompts() {
631+
return Flux.fromIterable(this.prompts.values())
632+
.map(McpStatelessServerFeatures.AsyncPromptSpecification::prompt);
633+
}
634+
616635
/**
617636
* Remove a prompt handler at runtime.
618637
* @param promptName The name of the prompt handler to remove
619638
* @return Mono that completes when clients have been notified of the change
620639
*/
621640
public Mono<Void> removePrompt(String promptName) {
622641
if (promptName == null) {
623-
return Mono.error(new McpError("Prompt name must not be null"));
642+
return Mono.error(new IllegalArgumentException("Prompt name must not be null"));
624643
}
625644
if (this.serverCapabilities.prompts() == null) {
626-
return Mono.error(new McpError("Server must be configured with prompt capabilities"));
645+
return Mono.error(new IllegalStateException("Server must be configured with prompt capabilities"));
627646
}
628647

629648
return Mono.defer(() -> {
@@ -633,7 +652,11 @@ public Mono<Void> removePrompt(String promptName) {
633652
logger.debug("Removed prompt handler: {}", promptName);
634653
return Mono.empty();
635654
}
636-
return Mono.error(new McpError("Prompt with name '" + promptName + "' not found"));
655+
else {
656+
logger.warn("Ignore as a Prompt with name '{}' not found", promptName);
657+
}
658+
659+
return Mono.empty();
637660
});
638661
}
639662

@@ -662,7 +685,10 @@ private McpStatelessRequestHandler<McpSchema.GetPromptResult> promptsGetRequestH
662685
// Implement prompt retrieval logic here
663686
McpStatelessServerFeatures.AsyncPromptSpecification specification = this.prompts.get(promptRequest.name());
664687
if (specification == null) {
665-
return Mono.error(new McpError("Prompt not found: " + promptRequest.name()));
688+
return Mono.error(McpError.builder(ErrorCodes.INVALID_PARAMS)
689+
.message("Invalid prompt name")
690+
.data("Prompt not found: " + promptRequest.name())
691+
.build());
666692
}
667693

668694
return specification.promptHandler().apply(ctx, promptRequest);

‎mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessSyncServer.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,14 @@ public void addTool(McpStatelessServerFeatures.SyncToolSpecification toolSpecifi
7474
.block();
7575
}
7676

77+
/**
78+
* List all registered tools.
79+
* @return A list of all registered tools
80+
*/
81+
public List<McpSchema.Tool> listTools() {
82+
return this.asyncServer.listTools().collectList().block();
83+
}
84+
7785
/**
7886
* Remove a tool handler at runtime.
7987
* @param toolName The name of the tool handler to remove
@@ -148,6 +156,14 @@ public void addPrompt(McpStatelessServerFeatures.SyncPromptSpecification promptS
148156
.block();
149157
}
150158

159+
/**
160+
* List all registered prompts.
161+
* @return A list of all registered prompts
162+
*/
163+
public List<McpSchema.Prompt> listPrompts() {
164+
return this.asyncServer.listPrompts().collectList().block();
165+
}
166+
151167
/**
152168
* Remove a prompt handler at runtime.
153169
* @param promptName The name of the prompt handler to remove

0 commit comments

Comments
 (0)