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:
parent
1e2a176c09
commit
94f172a44e
4 changed files with 153 additions and 26 deletions
|
|
@ -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());
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue