Skip to content

Commit 2748260

Browse files
tzolovtzolov
committed
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 modelcontextprotocol#46
1 parent c469f79 commit 2748260

File tree

8 files changed

+200
-98
lines changed

8 files changed

+200
-98
lines changed

‎mcp-docs/src/main/antora/modules/ROOT/images/MCP-layers.svg

Lines changed: 87 additions & 71 deletions
Loading

‎mcp-docs/src/main/antora/modules/ROOT/pages/mcp.adoc

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,13 @@ Add the following dependency to your Maven project:
3131

3232
== Architecture
3333

34-
The SDK follows a layered architecture with clear separation of concerns:
34+
image::mcp-layers.svg[width=400,float=right]
3535

36-
image::MCP-layers.svg[width=400,align=center]
36+
The SDK follows a layered architecture with clear separation of concerns:
3737

38-
* *Transport Layer (McpTransport)*: Handles JSON-RPC message serialization/deserialization via StdioTransport (stdin/stdout) and SseTransport (HTTP streaming).
39-
* *Session Layer (McpSession)*: Manages communication patterns and state using DefaultMcpSession implementation.
4038
* *Client/Server Layer*: Both use McpSession for sync/async operations, with McpClient handling client-side protocol operations and McpServer managing server-side protocol operations.
39+
* *Session Layer (McpSession)*: Manages communication patterns and state using DefaultMcpSession implementation.
40+
* *Transport Layer (McpTransport)*: Handles JSON-RPC message serialization/deserialization via StdioTransport (stdin/stdout) and SseTransport (HTTP streaming).
4141

4242
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.
4343

‎mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -409,7 +409,7 @@ private RequestHandler rootsListRequestHandler() {
409409

410410
List<Root> roots = this.roots.values().stream().toList();
411411

412-
return Mono.just(roots);
412+
return Mono.just(new McpSchema.ListRootsResult(roots));
413413
};
414414
}
415415

‎mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java

Lines changed: 26 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@
4040
import org.springframework.ai.mcp.spec.McpSchema;
4141
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
4242
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
43-
import org.springframework.ai.mcp.spec.McpSchema.ListRootsResult;
4443
import org.springframework.ai.mcp.spec.McpSchema.LoggingLevel;
4544
import org.springframework.ai.mcp.spec.McpSchema.LoggingMessageNotification;
4645
import org.springframework.ai.mcp.spec.McpSchema.Tool;
@@ -251,28 +250,35 @@ public void close() {
251250
private static TypeReference<McpSchema.ListRootsResult> LIST_ROOTS_RESULT_TYPE_REF = new TypeReference<>() {
252251
};
253252

254-
private NotificationHandler rootsListChnagedNotificationHandler(
255-
List<Consumer<List<McpSchema.Root>>> rootsChangeConsumers) {
256-
257-
if (this.clientCapabilities != null && this.clientCapabilities.roots() != null) {
253+
/**
254+
* Retrieves the list of all roots provided by the client.
255+
* @return A Mono that emits the list of roots result.
256+
*/
257+
public Mono<McpSchema.ListRootsResult> listRoots() {
258+
return this.listRoots(null);
259+
}
258260

259-
Mono<ListRootsResult> updatedRootsList = this.mcpSession.sendRequest(McpSchema.METHOD_ROOTS_LIST, null,
260-
LIST_ROOTS_RESULT_TYPE_REF);
261+
/**
262+
* Retrieves a paginated list of roots provided by the server.
263+
* @param cursor Optional pagination cursor from a previous list request
264+
* @return A Mono that emits the list of roots result containing
265+
*/
266+
public Mono<McpSchema.ListRootsResult> listRoots(String cursor) {
267+
return this.mcpSession.sendRequest(McpSchema.METHOD_ROOTS_LIST, new McpSchema.PaginatedRequest(cursor),
268+
LIST_ROOTS_RESULT_TYPE_REF);
269+
}
261270

262-
return params -> updatedRootsList // @formatter:off
263-
.flatMap(listRootsResult ->
264-
Mono.fromRunnable(() ->
265-
rootsChangeConsumers.stream().forEach(consumer -> consumer.accept(listRootsResult.roots())))
266-
.subscribeOn(Schedulers.boundedElastic())) // TODO: Check if this is needed
267-
.onErrorResume(error -> {
268-
logger.error("Error handling roots list change notification", error);
269-
return Mono.empty();
270-
})
271-
.then(); // Convert to Mono<Void>
272-
// @formatter:on
273-
}
271+
private NotificationHandler rootsListChnagedNotificationHandler(
272+
List<Consumer<List<McpSchema.Root>>> rootsChangeConsumers) {
274273

275-
return params -> Mono.empty();
274+
return params -> {
275+
return listRoots().flatMap(listRootsResult -> Mono.fromRunnable(() -> {
276+
rootsChangeConsumers.stream().forEach(consumer -> consumer.accept(listRootsResult.roots()));
277+
}).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> {
278+
logger.error("Error handling roots list change notification", error);
279+
return Mono.empty();
280+
}).then();
281+
};
276282
};
277283

278284
// ---------------------------------------

‎mcp/src/main/java/org/springframework/ai/mcp/server/McpSyncServer.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.springframework.ai.mcp.server.McpServer.PromptRegistration;
2020
import org.springframework.ai.mcp.server.McpServer.ResourceRegistration;
2121
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
22+
import org.springframework.ai.mcp.spec.McpError;
2223
import org.springframework.ai.mcp.spec.McpSchema;
2324
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
2425
import org.springframework.ai.mcp.spec.McpSchema.LoggingMessageNotification;
@@ -47,6 +48,23 @@ public McpSyncServer(McpAsyncServer asyncServer) {
4748
this.asyncServer = asyncServer;
4849
}
4950

51+
/**
52+
* Retrieves the list of all roots provided by the client.
53+
* @return The list of roots
54+
*/
55+
public McpSchema.ListRootsResult listRoots() {
56+
return this.listRoots(null);
57+
}
58+
59+
/**
60+
* Retrieves a paginated list of roots provided by the server.
61+
* @param cursor Optional pagination cursor from a previous list request
62+
* @return The list of roots
63+
*/
64+
public McpSchema.ListRootsResult listRoots(String cursor) {
65+
return this.asyncServer.listRoots(cursor).block();
66+
}
67+
5068
/**
5169
* Add a new tool handler.
5270
* @param toolHandler The tool handler to add

‎mcp/src/test/java/org/springframework/ai/mcp/MockMcpTransport.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,10 @@
3030
import org.springframework.ai.mcp.spec.McpSchema;
3131
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification;
3232
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
33+
import org.springframework.ai.mcp.spec.ServerMcpTransport;
3334

3435
@SuppressWarnings("unused")
35-
public class MockMcpTransport implements ClientMcpTransport {
36+
public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport {
3637

3738
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
3839

‎mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,8 @@ void testRootsListRequestHandling() {
100100

101101
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
102102
assertThat(response.id()).isEqualTo("test-id");
103-
assertThat(response.result()).isEqualTo(List.of(new Root("file:///test/path", "test-root")));
103+
assertThat(response.result())
104+
.isEqualTo(new McpSchema.ListRootsResult(List.of(new Root("file:///test/path", "test-root"))));
104105
assertThat(response.error()).isNull();
105106

106107
asyncMcpClient.closeGracefully();

‎mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@
1515
*/
1616
package org.springframework.ai.mcp.server;
1717

18+
import java.time.Duration;
1819
import java.util.List;
1920
import java.util.Map;
21+
import java.util.concurrent.atomic.AtomicReference;
2022
import java.util.function.Function;
2123

2224
import com.fasterxml.jackson.databind.ObjectMapper;
@@ -37,12 +39,14 @@
3739
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult;
3840
import org.springframework.ai.mcp.spec.McpSchema.InitializeResult;
3941
import org.springframework.ai.mcp.spec.McpSchema.Role;
42+
import org.springframework.ai.mcp.spec.McpSchema.Root;
4043
import org.springframework.http.server.reactive.HttpHandler;
4144
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
4245
import org.springframework.web.reactive.function.client.WebClient;
4346
import org.springframework.web.reactive.function.server.RouterFunctions;
4447

4548
import static org.assertj.core.api.Assertions.assertThat;
49+
import static org.awaitility.Awaitility.await;
4650

4751
public class SseAsyncIntegrationTests {
4852

@@ -156,4 +160,60 @@ void testCreateMessageSuccess() throws InterruptedException {
156160
}).verifyComplete();
157161
}
158162

163+
// ---------------------------------------
164+
// Roots Tests
165+
// ---------------------------------------
166+
@Test
167+
void testRootsSuccess() {
168+
169+
List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
170+
171+
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
172+
var mcpServer = McpServer.using(mcpServerTransport)
173+
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
174+
.sync();
175+
176+
// HttpHandler httpHandler =
177+
// RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction());
178+
// ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
179+
// HttpServer httpServer = HttpServer.create().port(8080).handle(adapter);
180+
// DisposableServer d = httpServer.bindNow();
181+
182+
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
183+
.roots(roots)
184+
.sync();
185+
186+
InitializeResult initResult = mcpClient.initialize();
187+
assertThat(initResult).isNotNull();
188+
189+
assertThat(rootsRef.get()).isNull();
190+
191+
assertThat(mcpServer.listRoots().roots()).containsAll(roots);
192+
193+
mcpClient.rootsListChangedNotification();
194+
195+
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
196+
assertThat(rootsRef.get()).containsAll(roots);
197+
});
198+
199+
// Remove a root
200+
mcpClient.removeRoot(roots.get(0).uri());
201+
202+
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
203+
assertThat(rootsRef.get()).containsAll(List.of(roots.get(1)));
204+
});
205+
206+
// Add a new root
207+
var root3 = new Root("uri3://", "root3");
208+
mcpClient.addRoot(root3);
209+
210+
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
211+
assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3));
212+
});
213+
214+
mcpClient.close();
215+
216+
mcpServer.close();
217+
}
218+
159219
}

0 commit comments

Comments
 (0)