From 2748260cdadb8dc1d5500dcab4e9e20df7235004 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sat, 4 Jan 2025 01:11:34 +0100 Subject: [PATCH] feat: Enhance MCP roots implementation and test coverage - Add proper ListRootsResult wrapping in async client - Implement paginated listRoots methods in server components - Fix roots notification handling in async server - Add roots integration tests - Update architecture diagram and documentation layout The changes improve the roots feature implementation with proper result types, and robust notification handling. Adds thorough test coverage for roots functionality including async notifications and root addition/removal scenarios. Resolves #46 --- .../antora/modules/ROOT/images/MCP-layers.svg | 166 ++++++++++-------- .../main/antora/modules/ROOT/pages/mcp.adoc | 8 +- .../ai/mcp/client/McpAsyncClient.java | 2 +- .../ai/mcp/server/McpAsyncServer.java | 46 ++--- .../ai/mcp/server/McpSyncServer.java | 18 ++ .../ai/mcp/MockMcpTransport.java | 3 +- .../McpAsyncClientResponseHandlerTests.java | 3 +- .../mcp/server/SseAsyncIntegrationTests.java | 60 +++++++ 8 files changed, 204 insertions(+), 102 deletions(-) diff --git a/mcp-docs/src/main/antora/modules/ROOT/images/MCP-layers.svg b/mcp-docs/src/main/antora/modules/ROOT/images/MCP-layers.svg index 2f38d6a..3847eaa 100644 --- a/mcp-docs/src/main/antora/modules/ROOT/images/MCP-layers.svg +++ b/mcp-docs/src/main/antora/modules/ROOT/images/MCP-layers.svg @@ -1,24 +1,24 @@ + id="defs27" /> + inkscape:current-layer="svg27"> + + + id="g27" + transform="translate(-334.47681,-179.33106)"> - + + - - + + diff --git a/mcp-docs/src/main/antora/modules/ROOT/pages/mcp.adoc b/mcp-docs/src/main/antora/modules/ROOT/pages/mcp.adoc index dc3755b..de98d5d 100644 --- a/mcp-docs/src/main/antora/modules/ROOT/pages/mcp.adoc +++ b/mcp-docs/src/main/antora/modules/ROOT/pages/mcp.adoc @@ -31,13 +31,13 @@ Add the following dependency to your Maven project: == Architecture +image::mcp-layers.svg[width=400,float=right] + The SDK follows a layered architecture with clear separation of concerns: -image::MCP-layers.svg[width=400,align=center] - -* *Transport Layer (McpTransport)*: Handles JSON-RPC message serialization/deserialization via StdioTransport (stdin/stdout) and SseTransport (HTTP streaming). -* *Session Layer (McpSession)*: Manages communication patterns and state using DefaultMcpSession implementation. * *Client/Server Layer*: Both use McpSession for sync/async operations, with McpClient handling client-side protocol operations and McpServer managing server-side protocol operations. +* *Session Layer (McpSession)*: Manages communication patterns and state using DefaultMcpSession implementation. +* *Transport Layer (McpTransport)*: Handles JSON-RPC message serialization/deserialization via StdioTransport (stdin/stdout) and SseTransport (HTTP streaming). Following class diagram illustrates the layered architecture of the MCP SDK, showing the relationships between core interfaces (McpTransport, McpSession), their implementations, and the client/server components. It highlights how the transport layer connects to sessions, which in turn support both synchronous and asynchronous client/server implementations. 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 e1bc8f1..ae9141f 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 @@ -409,7 +409,7 @@ public class McpAsyncClient { List roots = this.roots.values().stream().toList(); - return Mono.just(roots); + return Mono.just(new McpSchema.ListRootsResult(roots)); }; } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java index ca8cb83..ee6882d 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java @@ -40,7 +40,6 @@ import org.springframework.ai.mcp.spec.McpError; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.CallToolResult; import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities; -import org.springframework.ai.mcp.spec.McpSchema.ListRootsResult; import org.springframework.ai.mcp.spec.McpSchema.LoggingLevel; import org.springframework.ai.mcp.spec.McpSchema.LoggingMessageNotification; import org.springframework.ai.mcp.spec.McpSchema.Tool; @@ -251,28 +250,35 @@ public class McpAsyncServer { private static TypeReference LIST_ROOTS_RESULT_TYPE_REF = new TypeReference<>() { }; + /** + * Retrieves the list of all roots provided by the client. + * @return A Mono that emits the list of roots result. + */ + public Mono listRoots() { + return this.listRoots(null); + } + + /** + * Retrieves a paginated list of roots provided by the server. + * @param cursor Optional pagination cursor from a previous list request + * @return A Mono that emits the list of roots result containing + */ + public Mono listRoots(String cursor) { + return this.mcpSession.sendRequest(McpSchema.METHOD_ROOTS_LIST, new McpSchema.PaginatedRequest(cursor), + LIST_ROOTS_RESULT_TYPE_REF); + } + private NotificationHandler rootsListChnagedNotificationHandler( List>> rootsChangeConsumers) { - if (this.clientCapabilities != null && this.clientCapabilities.roots() != null) { - - Mono updatedRootsList = this.mcpSession.sendRequest(McpSchema.METHOD_ROOTS_LIST, null, - LIST_ROOTS_RESULT_TYPE_REF); - - return params -> updatedRootsList // @formatter:off - .flatMap(listRootsResult -> - Mono.fromRunnable(() -> - rootsChangeConsumers.stream().forEach(consumer -> consumer.accept(listRootsResult.roots()))) - .subscribeOn(Schedulers.boundedElastic())) // TODO: Check if this is needed - .onErrorResume(error -> { - logger.error("Error handling roots list change notification", error); - return Mono.empty(); - }) - .then(); // Convert to Mono - // @formatter:on - } - - return params -> Mono.empty(); + return params -> { + return listRoots().flatMap(listRootsResult -> Mono.fromRunnable(() -> { + rootsChangeConsumers.stream().forEach(consumer -> consumer.accept(listRootsResult.roots())); + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling roots list change notification", error); + return Mono.empty(); + }).then(); + }; }; // --------------------------------------- diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpSyncServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpSyncServer.java index 6096d9e..db564f0 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpSyncServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpSyncServer.java @@ -19,6 +19,7 @@ package org.springframework.ai.mcp.server; import org.springframework.ai.mcp.server.McpServer.PromptRegistration; import org.springframework.ai.mcp.server.McpServer.ResourceRegistration; import org.springframework.ai.mcp.server.McpServer.ToolRegistration; +import org.springframework.ai.mcp.spec.McpError; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities; import org.springframework.ai.mcp.spec.McpSchema.LoggingMessageNotification; @@ -47,6 +48,23 @@ public class McpSyncServer { this.asyncServer = asyncServer; } + /** + * Retrieves the list of all roots provided by the client. + * @return The list of roots + */ + public McpSchema.ListRootsResult listRoots() { + return this.listRoots(null); + } + + /** + * Retrieves a paginated list of roots provided by the server. + * @param cursor Optional pagination cursor from a previous list request + * @return The list of roots + */ + public McpSchema.ListRootsResult listRoots(String cursor) { + return this.asyncServer.listRoots(cursor).block(); + } + /** * Add a new tool handler. * @param toolHandler The tool handler to add diff --git a/mcp/src/test/java/org/springframework/ai/mcp/MockMcpTransport.java b/mcp/src/test/java/org/springframework/ai/mcp/MockMcpTransport.java index 034a08f..20736b1 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/MockMcpTransport.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/MockMcpTransport.java @@ -30,9 +30,10 @@ import org.springframework.ai.mcp.spec.ClientMcpTransport; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest; +import org.springframework.ai.mcp.spec.ServerMcpTransport; @SuppressWarnings("unused") -public class MockMcpTransport implements ClientMcpTransport { +public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport { private final AtomicInteger inboundMessageCount = new AtomicInteger(0); diff --git a/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java index da18e5f..81e0fd2 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java @@ -100,7 +100,8 @@ class McpAsyncClientResponseHandlerTests { McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage; assertThat(response.id()).isEqualTo("test-id"); - assertThat(response.result()).isEqualTo(List.of(new Root("file:///test/path", "test-root"))); + assertThat(response.result()) + .isEqualTo(new McpSchema.ListRootsResult(List.of(new Root("file:///test/path", "test-root")))); assertThat(response.error()).isNull(); asyncMcpClient.closeGracefully(); 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 65c7661..8e5caa5 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 @@ -15,8 +15,10 @@ */ package org.springframework.ai.mcp.server; +import java.time.Duration; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import com.fasterxml.jackson.databind.ObjectMapper; @@ -37,12 +39,14 @@ import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest; import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult; import org.springframework.ai.mcp.spec.McpSchema.InitializeResult; import org.springframework.ai.mcp.spec.McpSchema.Role; +import org.springframework.ai.mcp.spec.McpSchema.Root; import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; 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.awaitility.Awaitility.await; public class SseAsyncIntegrationTests { @@ -156,4 +160,60 @@ public class SseAsyncIntegrationTests { }).verifyComplete(); } + // --------------------------------------- + // Roots Tests + // --------------------------------------- + @Test + void testRootsSuccess() { + + List roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); + + AtomicReference> rootsRef = new AtomicReference<>(); + var mcpServer = McpServer.using(mcpServerTransport) + .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(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(rootsRef.get()).isNull(); + + assertThat(mcpServer.listRoots().roots()).containsAll(roots); + + mcpClient.rootsListChangedNotification(); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(roots); + }); + + // Remove a root + mcpClient.removeRoot(roots.get(0).uri()); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(List.of(roots.get(1))); + }); + + // Add a new root + var root3 = new Root("uri3://", "root3"); + mcpClient.addRoot(root3); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3)); + }); + + mcpClient.close(); + + mcpServer.close(); + } + }