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
4 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;
}
}

View file

@ -1,18 +1,18 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.mcp.server;
import java.time.Duration;
@ -46,6 +46,7 @@ import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.server.RouterFunctions;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.awaitility.Awaitility.await;
public class SseAsyncIntegrationTests {
@ -165,7 +166,6 @@ public class SseAsyncIntegrationTests {
// ---------------------------------------
@Test
void testRootsSuccess() {
List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
@ -173,12 +173,6 @@ public class SseAsyncIntegrationTests {
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync();
// HttpHandler httpHandler =
// RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction());
// ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
// HttpServer httpServer = HttpServer.create().port(8080).handle(adapter);
// DisposableServer d = httpServer.bindNow();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots)
.sync();
@ -212,8 +206,112 @@ public class SseAsyncIntegrationTests {
});
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsWithoutCapability() {
var mcpServer = McpServer.using(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> {
}).sync();
// Create client without roots capability
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()) // No
// roots
// capability
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
// Attempt to list roots should fail
assertThatThrownBy(() -> mcpServer.listRoots().roots()).isInstanceOf(McpError.class)
.hasMessage("Roots not supported");
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsWithEmptyRootsList() {
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(List.of()) // Empty roots list
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.rootsListChangedNotification();
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
assertThat(rootsRef.get()).isEmpty();
});
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsWithMultipleConsumers() {
List<Root> roots = List.of(new Root("uri1://", "root1"));
AtomicReference<List<Root>> rootsRef1 = new AtomicReference<>();
AtomicReference<List<Root>> rootsRef2 = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef1.set(rootsUpdate))
.rootsChangeConsumer(rootsUpdate -> rootsRef2.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots)
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.rootsListChangedNotification();
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
assertThat(rootsRef1.get()).containsAll(roots);
assertThat(rootsRef2.get()).containsAll(roots);
});
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsServerCloseWithActiveSubscription() {
List<Root> roots = List.of(new Root("uri1://", "root1"));
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots)
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.rootsListChangedNotification();
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
assertThat(rootsRef.get()).containsAll(roots);
});
// Close server while subscription is active
mcpServer.close();
// Verify client can handle server closure gracefully
mcpClient.close();
}
}