diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java b/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java index ae9141f..004cbc7 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java @@ -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 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()); } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java b/mcp/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java index 993b396..8a1dde0 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java @@ -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 diff --git a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpError.java b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpError.java index ad3e91f..94170ff 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpError.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpError.java @@ -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; + } + } \ No newline at end of file diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java index 8e5caa5..fb1f8aa 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java @@ -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 roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); AtomicReference> 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> 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 roots = List.of(new Root("uri1://", "root1")); + + AtomicReference> rootsRef1 = new AtomicReference<>(); + AtomicReference> 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 roots = List.of(new Root("uri1://", "root1")); + + AtomicReference> 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(); + } + }