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
This commit is contained in:
parent
c469f790b3
commit
2748260cda
3 changed files with 200 additions and 98 deletions
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
|
||||
|
||||
AtomicReference<List<Root>> 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();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue