Enhance roots functionality and improve error handling

- Replace Utils.isEmpty check with null check for roots map
- Add specialized error messages for method not found errors
- Extend McpError to include JSONRPCError information
- Add comprehensive tests for roots functionality including:
  - Empty roots list handling
  - Missing capability scenarios
  - Multiple consumers
  - Server closure with active subscriptions
This commit is contained in:
Christian Tzolov 2025-01-04 08:16:02 +01:00
commit 94f172a44e
3 changed files with 153 additions and 26 deletions

View file

@ -150,13 +150,13 @@ public class McpAsyncClient {
this.transport = transport;
this.roots = !Utils.isEmpty(roots) ? new ConcurrentHashMap<>(roots) : new ConcurrentHashMap<>();
this.roots = roots != null ? new ConcurrentHashMap<>(roots) : new ConcurrentHashMap<>();
// Request Handlers
Map<String, RequestHandler> requestHanlers = new HashMap<>();
// Roots List Request Handler
if (!Utils.isEmpty(this.roots) && this.clientCapabilities.roots() != null) {
if (this.roots != null && this.clientCapabilities.roots() != null) {
requestHanlers.put(McpSchema.METHOD_ROOTS_LIST, rootsListRequestHandler());
}

View file

@ -128,9 +128,11 @@ public class DefaultMcpSession implements McpSession {
this.requestHandlers.putAll(requestHandlers);
this.notificationHandlers.putAll(notificationHandlers);
// TODO: consider mono.transformDeferredContextual where the Context contains the
// TODO: consider mono.transformDeferredContextual where the Context contains
// the
// Observation associated with the individual message - it can be used to
// create child Observation and emit it together with the message to the consumer
// create child Observation and emit it together with the message to the
// consumer
this.connection = this.transport.connect(mono -> mono.doOnNext(message -> {
if (message instanceof McpSchema.JSONRPCResponse response) {
logger.info("Received Response: {}", response);
@ -169,9 +171,10 @@ public class DefaultMcpSession implements McpSession {
return Mono.defer(() -> {
var handler = this.requestHandlers.get(request.method());
if (handler == null) {
MethodNotFoundError error = getMethodNotFoundError(request.method());
return Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND,
"Method not found: " + request.method(), null)));
error.message(), error.data())));
}
return handler.handle(request.params())
@ -183,6 +186,19 @@ public class DefaultMcpSession implements McpSession {
});
}
record MethodNotFoundError(String method, String message, Object data) {
}
public static MethodNotFoundError getMethodNotFoundError(String method) {
switch (method) {
case McpSchema.METHOD_ROOTS_LIST:
return new MethodNotFoundError(method, "Roots not supported",
Map.of("reason", "Client does not have roots capability"));
default:
return new MethodNotFoundError(method, "Method not found: " + method, null);
}
}
/**
* Handles an incoming JSON-RPC notification by routing it to the appropriate handler.
* @param notification The incoming JSON-RPC notification

View file

@ -15,10 +15,23 @@
*/
package org.springframework.ai.mcp.spec;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCResponse.JSONRPCError;
public class McpError extends RuntimeException {
private JSONRPCError jsonRpcError;
public McpError(JSONRPCError jsonRpcError) {
super(jsonRpcError.message());
this.jsonRpcError = jsonRpcError;
}
public McpError(Object error) {
super(error.toString());
}
public JSONRPCError getJsonRpcError() {
return jsonRpcError;
}
}