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.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());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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;
|
||||||
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
@ -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();
|
||||||
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue