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.transport = transport;
this.roots = !Utils.isEmpty(roots) ? new ConcurrentHashMap<>(roots) : new ConcurrentHashMap<>(); this.roots = roots != null ? new ConcurrentHashMap<>(roots) : new ConcurrentHashMap<>();
// Request Handlers // Request Handlers
Map<String, RequestHandler> requestHanlers = new HashMap<>(); Map<String, RequestHandler> requestHanlers = new HashMap<>();
// Roots List Request Handler // 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()); requestHanlers.put(McpSchema.METHOD_ROOTS_LIST, rootsListRequestHandler());
} }

View file

@ -128,9 +128,11 @@ public class DefaultMcpSession implements McpSession {
this.requestHandlers.putAll(requestHandlers); this.requestHandlers.putAll(requestHandlers);
this.notificationHandlers.putAll(notificationHandlers); 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 // 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 -> { this.connection = this.transport.connect(mono -> mono.doOnNext(message -> {
if (message instanceof McpSchema.JSONRPCResponse response) { if (message instanceof McpSchema.JSONRPCResponse response) {
logger.info("Received Response: {}", response); logger.info("Received Response: {}", response);
@ -169,9 +171,10 @@ public class DefaultMcpSession implements McpSession {
return Mono.defer(() -> { return Mono.defer(() -> {
var handler = this.requestHandlers.get(request.method()); var handler = this.requestHandlers.get(request.method());
if (handler == null) { if (handler == null) {
MethodNotFoundError error = getMethodNotFoundError(request.method());
return Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, return Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND, new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND,
"Method not found: " + request.method(), null))); error.message(), error.data())));
} }
return handler.handle(request.params()) 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. * Handles an incoming JSON-RPC notification by routing it to the appropriate handler.
* @param notification The incoming JSON-RPC notification * @param notification The incoming JSON-RPC notification

View file

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

View file

@ -1,18 +1,18 @@
/* /*
* Copyright 2024 - 2024 the original author or authors. * Copyright 2024 - 2024 the original author or authors.
* *
* Licensed under the Apache License, Version 2.0 (the "License"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. * you may not use this file except in compliance with the License.
* You may obtain a copy of the License at * You may obtain a copy of the License at
* *
* https://www.apache.org/licenses/LICENSE-2.0 * https://www.apache.org/licenses/LICENSE-2.0
* *
* Unless required by applicable law or agreed to in writing, software * Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, * distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and * See the License for the specific language governing permissions and
* limitations under the License. * limitations under the License.
*/ */
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import java.time.Duration; 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 org.springframework.web.reactive.function.server.RouterFunctions;
import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.awaitility.Awaitility.await; import static org.awaitility.Awaitility.await;
public class SseAsyncIntegrationTests { public class SseAsyncIntegrationTests {
@ -165,7 +166,6 @@ public class SseAsyncIntegrationTests {
// --------------------------------------- // ---------------------------------------
@Test @Test
void testRootsSuccess() { void testRootsSuccess() {
List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
AtomicReference<List<Root>> rootsRef = new AtomicReference<>(); AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
@ -173,12 +173,6 @@ public class SseAsyncIntegrationTests {
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync(); .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()) var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots) .roots(roots)
.sync(); .sync();
@ -212,8 +206,112 @@ public class SseAsyncIntegrationTests {
}); });
mcpClient.close(); mcpClient.close();
mcpServer.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();
}
} }