feat: Add roots change notification support to MCP server

- Implement roots change notification handler in McpAsyncServer
- Add rootsChangeConsumer builder methods to support single and multiple consumers
- Rename info() builder method to serverInfo() for clarity
- Add comprehensive tests for roots change notification functionality
- Add default server info implementation in Builder

Resolves #39
This commit is contained in:
Christian Tzolov 2025-01-02 08:27:12 +01:00
commit 0fd93a9f82
2 changed files with 266 additions and 58 deletions

View file

@ -23,6 +23,7 @@ import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Consumer;
import com.fasterxml.jackson.core.type.TypeReference;
import org.slf4j.Logger;
@ -39,6 +40,7 @@ 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;
@ -84,7 +86,6 @@ public class McpAsyncServer {
private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG;
// TODO: Add support for roots list changed notification
/**
* Create a new McpAsyncServer with the given transport and capabilities.
* @param mcpTransport The transport layer implementation for MCP communication
@ -94,11 +95,13 @@ public class McpAsyncServer {
* @param resources The map of resource registrations
* @param resourceTemplates The list of resource templates
* @param prompts The map of prompt registrations
* @param rootsChangeConsumers The list of consumers that will be notified when the
* roots list changes
*/
public McpAsyncServer(McpTransport mcpTransport, McpSchema.Implementation serverInfo,
McpSchema.ServerCapabilities serverCapabilities, List<ToolRegistration> tools,
Map<String, ResourceRegistration> resources, List<McpSchema.ResourceTemplate> resourceTemplates,
Map<String, PromptRegistration> prompts) {
Map<String, PromptRegistration> prompts, List<Consumer<List<McpSchema.Root>>> rootsChangeConsumers) {
this.serverInfo = serverInfo;
this.tools = new CopyOnWriteArrayList<>(tools != null ? tools : List.of());
@ -156,6 +159,13 @@ public class McpAsyncServer {
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, (params) -> Mono.empty());
if (Utils.isEmpty(rootsChangeConsumers)) {
rootsChangeConsumers = List.of((roots) -> logger
.warn("Roots list changed notification, but no consumers provided. Roots list changed: {}", roots));
}
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED,
rootsListChnagedNotificationHandler(rootsChangeConsumers));
this.transport = mcpTransport;
this.mcpSession = new DefaultMcpSession(Duration.ofSeconds(10), mcpTransport, requestHandlers,
notificationHandlers);
@ -238,6 +248,33 @@ public class McpAsyncServer {
this.mcpSession.close();
}
private static TypeReference<McpSchema.ListRootsResult> LIST_ROOTS_RESULT_TYPE_REF = new TypeReference<>() {
};
private NotificationHandler rootsListChnagedNotificationHandler(
List<Consumer<List<McpSchema.Root>>> rootsChangeConsumers) {
if (this.clientCapabilities != null && this.clientCapabilities.roots() != null) {
Mono<ListRootsResult> 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<Void>
// @formatter:on
}
return params -> Mono.empty();
};
// ---------------------------------------
// Tool Management
// ---------------------------------------

View file

@ -20,6 +20,7 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;
import org.springframework.ai.mcp.spec.McpSchema;
@ -206,9 +207,12 @@ public interface McpServer {
*/
public static class Builder {
private final static McpSchema.Implementation DEFAULT_SERVER_INFO = new McpSchema.Implementation("mcp-server",
"1.0.0");
private final McpTransport transport;
private McpSchema.Implementation serverInfo;
private McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO;
private McpSchema.ServerCapabilities serverCapabilities;
@ -241,6 +245,8 @@ public interface McpServer {
*/
private Map<String, PromptRegistration> prompts = new HashMap<>();
private List<Consumer<List<McpSchema.Root>>> rootsChangeConsumers = new ArrayList<>();
private Builder(McpTransport transport) {
Assert.notNull(transport, "Transport must not be null");
this.transport = transport;
@ -255,7 +261,7 @@ public interface McpServer {
* @return This builder instance for method chaining
* @throws IllegalArgumentException if serverInfo is null
*/
public Builder info(McpSchema.Implementation serverInfo) {
public Builder serverInfo(McpSchema.Implementation serverInfo) {
Assert.notNull(serverInfo, "Server info must not be null");
this.serverInfo = serverInfo;
return this;
@ -263,14 +269,15 @@ public interface McpServer {
/**
* Sets the server implementation information using name and version strings. This
* is a convenience method alternative to {@link #info(McpSchema.Implementation)}.
* is a convenience method alternative to
* {@link #serverInfo(McpSchema.Implementation)}.
* @param name The server name. Must not be null or empty.
* @param version The server version. Must not be null or empty.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if name or version is null or empty
* @see #info(McpSchema.Implementation)
* @see #serverInfo(McpSchema.Implementation)
*/
public Builder info(String name, String version) {
public Builder serverInfo(String name, String version) {
Assert.hasText(name, "Name must not be null or empty");
Assert.hasText(version, "Version must not be null or empty");
this.serverInfo = new McpSchema.Implementation(name, version);
@ -518,6 +525,49 @@ public interface McpServer {
return this;
}
/**
* Registers a consumer that will be notified when the list of roots changes. This
* is useful for updating resource availability dynamically, such as when new
* files are added or removed.
* @param consumer The consumer to register. Must not be null.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if consumer is null
*/
public Builder rootsChangeConsumer(Consumer<List<McpSchema.Root>> consumer) {
Assert.notNull(consumer, "Consumer must not be null");
this.rootsChangeConsumers.add(consumer);
return this;
}
/**
* Registers multiple consumers that will be notified when the list of roots
* changes. This method is useful when multiple consumers need to be registered at
* once.
* @param consumers The list of consumers to register. Must not be null.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if consumers is null
*/
public Builder rootsChangeConsumers(List<Consumer<List<McpSchema.Root>>> consumers) {
Assert.notNull(consumers, "Consumers list must not be null");
this.rootsChangeConsumers.addAll(consumers);
return this;
}
/**
* Registers multiple consumers that will be notified when the list of roots
* changes using varargs. This method provides a convenient way to register
* multiple consumers inline.
* @param consumers The consumers to register. Must not be null.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if consumers is null
*/
public Builder rootsChangeConsumers(Consumer<List<McpSchema.Root>>... consumers) {
for (Consumer<List<McpSchema.Root>> consumer : consumers) {
this.rootsChangeConsumers.add(consumer);
}
return this;
}
/**
* Builds a synchronous MCP server that provides blocking operations. Synchronous
* servers process each request to completion before handling the next one, making
@ -539,12 +589,9 @@ public interface McpServer {
* settings
*/
public McpAsyncServer async() {
if (serverInfo == null) {
serverInfo = new McpSchema.Implementation("mcp-server", "1.0.0");
}
return new McpAsyncServer(transport, serverInfo, serverCapabilities, tools, resources, resourceTemplates,
prompts);
prompts, rootsChangeConsumers);
}
}