From 383d6a7effb328053830a4c982e5ce86ceed0402 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 19 Dec 2024 13:51:30 +0100 Subject: [PATCH] feat(mcp-client): Add non-blocking tools change notification support - Add toolsChangeNotificationHandler with non-blocking consumer execution using boundedElastic scheduler - Add toolsChangeConsumer builder method to McpClient for easy configuration - Add comprehensive test coverage for tools change notifications - Update documentation with tools change notification features and examples The implementation ensures consumers are executed non-blockingly on the boundedElastic scheduler, preventing performance impact from blocking implementations while maintaining proper error handling and logging. Part of #20 --- spring-ai-mcp-core/README.md | 61 ++++++++++++++++- spring-ai-mcp-core/pom.xml | 9 ++- .../ai/mcp/client/McpAsyncClient.java | 33 ++++++++- .../ai/mcp/client/McpClient.java | 12 +++- .../client/AbstractMcpAsyncClientTests.java | 2 +- .../McpAsyncClientResponseHandlerTests.java | 68 ++++++++++++++----- 6 files changed, 162 insertions(+), 23 deletions(-) diff --git a/spring-ai-mcp-core/README.md b/spring-ai-mcp-core/README.md index 60a092f..8f259cf 100644 --- a/spring-ai-mcp-core/README.md +++ b/spring-ai-mcp-core/README.md @@ -12,7 +12,8 @@ This SDK implements the Model Context Protocol, enabling seamless integration wi - Standard MCP operations support: - Protocol version compatibility negotiation - Client-server capability exchange - - Tool discovery and execution + - Tool discovery and execution with change notifications + - Tool list change notifications with non-blocking consumer support - Resource management with URI templates - Resource subscription system - Roots list management and notifications @@ -185,6 +186,23 @@ var subscription = client.listResources() return client.subscribeResource(new McpSchema.SubscribeRequest("resource-uri")); }); +// Set up tools change notification handling +List>> toolsChangeConsumers = List.of( + tools -> { + // Handle tools list changes reactively + tools.forEach(tool -> { + System.out.println("Tool updated: " + tool.name()); + }); + } +); + +McpAsyncClient clientWithToolsNotifications = McpClient.using(transport) + .toolsChangeConsumer(toolsChangeConsumers) + .async(); + +// The client will now automatically handle tools/list_changed notifications +// and invoke the consumers on the boundedElastic scheduler to avoid blocking + // Handle results reactively or block if needed McpSchema.GetPromptResult promptResult = result.block(); subscription.block(); @@ -243,6 +261,47 @@ The SDK follows a layered architecture with clear separation of concerns: - Execution handling with timeout support - Result processing with error handling +### Tool Change Notifications + +The SDK supports automatic handling of tool list changes through a non-blocking notification system: + +#### Features +- Register multiple consumers to handle tool list changes +- Non-blocking execution using Project Reactor's boundedElastic scheduler +- Automatic tools/list request handling when notifications are received +- Error resilient with proper error handling and logging + +#### Example with Tools Change Notification + +```java +// Create tool change consumers +List>> toolsChangeConsumers = List.of( + tools -> { + // First consumer - e.g., update UI + tools.forEach(tool -> updateToolsUI(tool)); + }, + tools -> { + // Second consumer - e.g., update cache + toolsCache.updateTools(tools); + } +); + +// Create client with tools change notification support +McpAsyncClient client = McpClient.using(transport) + .toolsChangeConsumer(toolsChangeConsumers) + .async(); + +// Initialize client +client.initialize() + .doOnSuccess(result -> { + // Client will automatically handle tools/list_changed notifications + // and invoke consumers non-blockingly on boundedElastic scheduler + }) + .subscribe(); +``` + +The tools change notification system ensures that consumers are executed non-blockingly, preventing any potential performance impact from blocking implementations. All consumers are executed on Project Reactor's boundedElastic scheduler, making it safe to perform potentially blocking operations within the consumers. + ## Error Handling The SDK provides comprehensive error handling through the McpError class: diff --git a/spring-ai-mcp-core/pom.xml b/spring-ai-mcp-core/pom.xml index 4e143cd..9614718 100644 --- a/spring-ai-mcp-core/pom.xml +++ b/spring-ai-mcp-core/pom.xml @@ -98,7 +98,14 @@ test + + org.awaitility + awaitility + 4.2.0 + test + + - \ No newline at end of file + diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java index f5e55a4..d9911bb 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java @@ -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.Supplier; import com.fasterxml.jackson.core.type.TypeReference; @@ -29,6 +30,7 @@ import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; import org.springframework.ai.mcp.spec.DefaultMcpSession; +import org.springframework.ai.mcp.spec.DefaultMcpSession.NotificationHandler; import org.springframework.ai.mcp.spec.DefaultMcpSession.RequestHandler; import org.springframework.ai.mcp.spec.McpError; import org.springframework.ai.mcp.spec.McpSchema; @@ -78,7 +80,8 @@ public class McpAsyncClient { * notification. */ public McpAsyncClient(McpTransport transport, Duration requestTimeout, - List>> rootsListProviders, boolean rootsListChangedNotification) { + List>> rootsListProviders, boolean rootsListChangedNotification, + List>> toolsChangeConsumers) { Map requestHanlers = new HashMap<>(); @@ -87,7 +90,14 @@ public class McpAsyncClient { this.rootCapabilities = new McpSchema.ClientCapabilities.RootCapabilities(rootsListChangedNotification); } - this.mcpSession = new DefaultMcpSession(requestTimeout, transport, requestHanlers, Map.of()); + Map notificationHandlers = new HashMap<>(); + + if (toolsChangeConsumers != null && !toolsChangeConsumers.isEmpty()) { + notificationHandlers.put("notifications/tools/list_changed", + toolsChangeNotificationHandler(toolsChangeConsumers)); + } + + this.mcpSession = new DefaultMcpSession(requestTimeout, transport, requestHanlers, notificationHandlers); } @@ -116,6 +126,25 @@ public class McpAsyncClient { }; }; + private NotificationHandler toolsChangeNotificationHandler( + List>> toolsChangeConsumers) { + + return new NotificationHandler() { + @Override + public Mono handle(Object params) { + // TODO: add support for cursor/pagination + return listTools().flatMap(listToolsResult -> Mono.fromRunnable(() -> { + for (Consumer> toolsChangeConsumer : toolsChangeConsumers) { + toolsChangeConsumer.accept(listToolsResult.tools()); + } + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling tools list change notification", error); + return Mono.empty(); + }).then(); // Convert to Mono + } + }; + }; + /** * The initialization phase MUST be the first interaction between client and server. * During this phase, the client and server: diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java index 847849e..fb579b1 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java @@ -19,8 +19,10 @@ package org.springframework.ai.mcp.client; import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.function.Consumer; import java.util.function.Supplier; +import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.Root; import org.springframework.ai.mcp.spec.McpTransport; import org.springframework.ai.mcp.util.Assert; @@ -84,6 +86,8 @@ public class McpClient { private List>> rootsListProviders = new ArrayList<>(); + private List>> toolsChangeConsumers = new ArrayList<>(); + private Builder(McpTransport transport) { Assert.notNull(transport, "Transport must not be null"); this.transport = transport; @@ -110,6 +114,11 @@ public class McpClient { return this; } + public Builder toolsChangeConsumer(Consumer> toolsChangeConsumer) { + this.toolsChangeConsumers.add(toolsChangeConsumer); + return this; + } + /** * Build a synchronous MCP client. * @return A new instance of {@link McpSyncClient} @@ -123,7 +132,8 @@ public class McpClient { * @return A new instance of {@link McpAsyncClient} */ public McpAsyncClient async() { - return new McpAsyncClient(transport, requestTimeout, rootsListProviders, rootsListChangedNotification); + return new McpAsyncClient(transport, requestTimeout, rootsListProviders, rootsListChangedNotification, + toolsChangeConsumers); } } diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java index eba0302..c8cc83f 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java @@ -183,7 +183,7 @@ public abstract class AbstractMcpAsyncClientTests { var transport = createMcpTransport(); List>> providers = List.of(() -> List.of(new Root("file:///test/path", "test-root"))); - var client = new McpAsyncClient(transport, TIMEOUT, providers, true); + var client = new McpAsyncClient(transport, TIMEOUT, providers, true, List.of()); assertThatCode(() -> client.initialize().block(Duration.ofSeconds(10))).doesNotThrowAnyException(); diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java index c72adcd..02777fc 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java @@ -16,15 +16,19 @@ package org.springframework.ai.mcp.client; +import java.time.Duration; +import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; import com.fasterxml.jackson.core.type.TypeReference; + +import static org.awaitility.Awaitility.await; import com.fasterxml.jackson.databind.ObjectMapper; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -41,10 +45,6 @@ import static org.assertj.core.api.Assertions.assertThat; class McpAsyncClientResponseHandlerTests { - private McpAsyncClient asyncMcpClient; - - private MockMcpTransport transport; - @SuppressWarnings("unused") private static class MockMcpTransport implements McpTransport { @@ -103,24 +103,56 @@ class McpAsyncClientResponseHandlerTests { } - @BeforeEach - void setUp() { - transport = new MockMcpTransport(); + @Test + void testToolsChangeNotificationHandling() { + MockMcpTransport transport = new MockMcpTransport(); - Supplier> rootsListProvider = () -> List.of(new Root("file:///test/path", "test-root")); + // Create a list to store received tools for verification + List receivedTools = new ArrayList<>(); - asyncMcpClient = McpClient.using(transport).rootsListProvider(rootsListProvider).async(); - } + // Create a consumer that will be called when tools change + Consumer> toolsChangeConsumer = tools -> { + receivedTools.addAll(tools); + }; - @AfterEach - void tearDown() { - if (asyncMcpClient != null) { - asyncMcpClient.closeGracefully(); - } + // Create client with tools change consumer + McpAsyncClient asyncMcpClient = McpClient.using(transport).toolsChangeConsumer(toolsChangeConsumer).async(); + + // Create a mock tools list that the server will return + Map inputSchema = Map.of("type", "object", "properties", Map.of(), "required", List.of()); + McpSchema.Tool mockTool = new McpSchema.Tool("test-tool", "Test Tool Description", inputSchema); + McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(mockTool), null); + + // Simulate server sending tools/list_changed notification + McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, + "notifications/tools/list_changed", null); + transport.simulateIncomingMessage(notification); + + // Simulate server response to tools/list request + McpSchema.JSONRPCRequest toolsListRequest = transport.getLastSentMessageAsRequest(); + assertThat(toolsListRequest.method()).isEqualTo("tools/list"); + + McpSchema.JSONRPCResponse toolsListResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + toolsListRequest.id(), mockToolsResult, null); + transport.simulateIncomingMessage(toolsListResponse); + + // Verify the consumer received the expected tools + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(receivedTools).hasSize(1); + assertThat(receivedTools.get(0).name()).isEqualTo("test-tool"); + assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool Description"); + }); + + asyncMcpClient.closeGracefully(); } @Test void testRootsListRequestHandling() { + MockMcpTransport transport = new MockMcpTransport(); + + Supplier> rootsListProvider = () -> List.of(new Root("file:///test/path", "test-root")); + + McpAsyncClient asyncMcpClient = McpClient.using(transport).rootsListProvider(rootsListProvider).async(); // Simulate incoming request McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "roots/list", @@ -135,6 +167,8 @@ class McpAsyncClientResponseHandlerTests { assertThat(response.id()).isEqualTo("test-id"); assertThat(response.result()).isEqualTo(List.of(new Root("file:///test/path", "test-root"))); assertThat(response.error()).isNull(); + + asyncMcpClient.closeGracefully(); } }