From 975578b64b2372797942101901c3ffbf30e79736 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 19 Dec 2024 14:23:55 +0100 Subject: [PATCH] feat: Add resources and prompts change notification support - Add resourcesChangeNotificationHandler and promptsChangeNotificationHandler with non-blocking execution - Add builder methods for resources and prompts change consumers - Add test coverage for all notification handlers - Update documentation to cover all three notification types (tools, resources, prompts) - Consolidate notification documentation under unified Change Notifications section - Ensure consistent non-blocking behavior using boundedElastic scheduler All notification handlers follow the same pattern of non-blocking execution and error handling, providing a consistent approach to change notifications across the system. Resolves #20 --- spring-ai-mcp-core/README.md | 47 ++++++++-- .../ai/mcp/client/McpAsyncClient.java | 52 ++++++++++- .../ai/mcp/client/McpClient.java | 16 +++- .../client/AbstractMcpAsyncClientTests.java | 2 +- .../McpAsyncClientResponseHandlerTests.java | 92 +++++++++++++++++++ 5 files changed, 197 insertions(+), 12 deletions(-) diff --git a/spring-ai-mcp-core/README.md b/spring-ai-mcp-core/README.md index 8f259cf..3da5e59 100644 --- a/spring-ai-mcp-core/README.md +++ b/spring-ai-mcp-core/README.md @@ -261,17 +261,17 @@ The SDK follows a layered architecture with clear separation of concerns: - Execution handling with timeout support - Result processing with error handling -### Tool Change Notifications +### Change Notifications -The SDK supports automatic handling of tool list changes through a non-blocking notification system: +The SDK supports automatic handling of changes through a non-blocking notification system for tools, resources, and prompts: #### Features -- Register multiple consumers to handle tool list changes +- Register multiple consumers for tools, resources, and prompts changes - Non-blocking execution using Project Reactor's boundedElastic scheduler -- Automatic tools/list request handling when notifications are received +- Automatic list request handling when notifications are received - Error resilient with proper error handling and logging -#### Example with Tools Change Notification +#### Example with Change Notifications ```java // Create tool change consumers @@ -286,21 +286,50 @@ List>> toolsChangeConsumers = List.of( } ); -// Create client with tools change notification support +// Create resource change consumers +List>> resourcesChangeConsumers = List.of( + resources -> { + // Handle resource changes + resources.forEach(resource -> { + System.out.println("Resource updated: " + resource.uri()); + updateResourcesUI(resource); + }); + } +); + +// Create prompt change consumers +List>> promptsChangeConsumers = List.of( + prompts -> { + // Handle prompt changes + prompts.forEach(prompt -> { + System.out.println("Prompt updated: " + prompt.name()); + updatePromptsCache(prompt); + }); + } +); + +// Create client with change notification support McpAsyncClient client = McpClient.using(transport) - .toolsChangeConsumer(toolsChangeConsumers) + .toolsChangeConsumer(toolsChangeConsumer) + .resourcesChangeConsumer(resourcesChangeConsumer) + .promptsChangeConsumer(promptsChangeConsumer) .async(); // Initialize client client.initialize() .doOnSuccess(result -> { - // Client will automatically handle tools/list_changed notifications + // Client will automatically handle all change 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. +The change notification system ensures that all 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. + +Each type of notification handler operates independently: +- Tools: Handles notifications when available tools change +- Resources: Handles notifications when available resources change +- Prompts: Handles notifications when available prompts change ## Error Handling 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 d9911bb..ba3346e 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 @@ -81,7 +81,9 @@ public class McpAsyncClient { */ public McpAsyncClient(McpTransport transport, Duration requestTimeout, List>> rootsListProviders, boolean rootsListChangedNotification, - List>> toolsChangeConsumers) { + List>> toolsChangeConsumers, + List>> resourcesChangeConsumers, + List>> promptsChangeConsumers) { Map requestHanlers = new HashMap<>(); @@ -97,6 +99,16 @@ public class McpAsyncClient { toolsChangeNotificationHandler(toolsChangeConsumers)); } + if (resourcesChangeConsumers != null && !resourcesChangeConsumers.isEmpty()) { + notificationHandlers.put("notifications/resources/list_changed", + resourcesChangeNotificationHandler(resourcesChangeConsumers)); + } + + if (promptsChangeConsumers != null && !promptsChangeConsumers.isEmpty()) { + notificationHandlers.put("notifications/prompts/list_changed", + promptsChangeNotificationHandler(promptsChangeConsumers)); + } + this.mcpSession = new DefaultMcpSession(requestTimeout, transport, requestHanlers, notificationHandlers); } @@ -145,6 +157,44 @@ public class McpAsyncClient { }; }; + private NotificationHandler resourcesChangeNotificationHandler( + List>> resourcesChangeConsumers) { + + return new NotificationHandler() { + @Override + public Mono handle(Object params) { + // TODO: add support for cursor/pagination + return listResources().flatMap(listResourcesResult -> Mono.fromRunnable(() -> { + for (Consumer> resourceChangeConsumer : resourcesChangeConsumers) { + resourceChangeConsumer.accept(listResourcesResult.resources()); + } + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling resources list change notification", error); + return Mono.empty(); + }).then(); // Convert to Mono + } + }; + }; + + private NotificationHandler promptsChangeNotificationHandler( + List>> promptsChangeConsumers) { + + return new NotificationHandler() { + @Override + public Mono handle(Object params) { + // TODO: add support for cursor/pagination + return listPrompts().flatMap(listPromptsResult -> Mono.fromRunnable(() -> { + for (Consumer> promptChangeConsumer : promptsChangeConsumers) { + promptChangeConsumer.accept(listPromptsResult.prompts()); + } + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling prompts 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 fb579b1..04375ee 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 @@ -88,6 +88,10 @@ public class McpClient { private List>> toolsChangeConsumers = new ArrayList<>(); + private List>> resourcesChangeConsumers = new ArrayList<>(); + + private List>> promptsChangeConsumers = new ArrayList<>(); + private Builder(McpTransport transport) { Assert.notNull(transport, "Transport must not be null"); this.transport = transport; @@ -119,6 +123,16 @@ public class McpClient { return this; } + public Builder resourcesChangeConsumer(Consumer> resourcesChangeConsumer) { + this.resourcesChangeConsumers.add(resourcesChangeConsumer); + return this; + } + + public Builder promptsChangeConsumer(Consumer> promptsChangeConsumer) { + this.promptsChangeConsumers.add(promptsChangeConsumer); + return this; + } + /** * Build a synchronous MCP client. * @return A new instance of {@link McpSyncClient} @@ -133,7 +147,7 @@ public class McpClient { */ public McpAsyncClient async() { return new McpAsyncClient(transport, requestTimeout, rootsListProviders, rootsListChangedNotification, - toolsChangeConsumers); + toolsChangeConsumers, resourcesChangeConsumers, promptsChangeConsumers); } } 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 c8cc83f..ac2bb53 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, List.of()); + var client = McpClient.using(transport).requestTimeout(TIMEOUT).rootsListProvider(providers.get(0)).async(); 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 02777fc..e172f92 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 @@ -171,4 +171,96 @@ class McpAsyncClientResponseHandlerTests { asyncMcpClient.closeGracefully(); } + @Test + void testResourcesChangeNotificationHandling() { + MockMcpTransport transport = new MockMcpTransport(); + + // Create a list to store received resources for verification + List receivedResources = new ArrayList<>(); + + // Create a consumer that will be called when resources change + Consumer> resourcesChangeConsumer = resources -> { + receivedResources.addAll(resources); + }; + + // Create client with resources change consumer + McpAsyncClient asyncMcpClient = McpClient.using(transport) + .resourcesChangeConsumer(resourcesChangeConsumer) + .async(); + + // Create a mock resources list that the server will return + McpSchema.Resource mockResource = new McpSchema.Resource("test://resource", "Test Resource", "A test resource", + "text/plain", null); + McpSchema.ListResourcesResult mockResourcesResult = new McpSchema.ListResourcesResult(List.of(mockResource), + null); + + // Simulate server sending resources/list_changed notification + McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, + "notifications/resources/list_changed", null); + transport.simulateIncomingMessage(notification); + + // Simulate server response to resources/list request + McpSchema.JSONRPCRequest resourcesListRequest = transport.getLastSentMessageAsRequest(); + assertThat(resourcesListRequest.method()).isEqualTo("resources/list"); + + McpSchema.JSONRPCResponse resourcesListResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + resourcesListRequest.id(), mockResourcesResult, null); + transport.simulateIncomingMessage(resourcesListResponse); + + // Verify the consumer received the expected resources + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(receivedResources).hasSize(1); + assertThat(receivedResources.get(0).uri()).isEqualTo("test://resource"); + assertThat(receivedResources.get(0).name()).isEqualTo("Test Resource"); + assertThat(receivedResources.get(0).description()).isEqualTo("A test resource"); + }); + + asyncMcpClient.closeGracefully(); + } + + @Test + void testPromptsChangeNotificationHandling() { + MockMcpTransport transport = new MockMcpTransport(); + + // Create a list to store received prompts for verification + List receivedPrompts = new ArrayList<>(); + + // Create a consumer that will be called when prompts change + Consumer> promptsChangeConsumer = prompts -> { + receivedPrompts.addAll(prompts); + }; + + // Create client with prompts change consumer + McpAsyncClient asyncMcpClient = McpClient.using(transport).promptsChangeConsumer(promptsChangeConsumer).async(); + + // Create a mock prompts list that the server will return + McpSchema.Prompt mockPrompt = new McpSchema.Prompt("test-prompt", "Test Prompt Description", + List.of(new McpSchema.PromptArgument("arg1", "Test argument", true))); + McpSchema.ListPromptsResult mockPromptsResult = new McpSchema.ListPromptsResult(List.of(mockPrompt), null); + + // Simulate server sending prompts/list_changed notification + McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, + "notifications/prompts/list_changed", null); + transport.simulateIncomingMessage(notification); + + // Simulate server response to prompts/list request + McpSchema.JSONRPCRequest promptsListRequest = transport.getLastSentMessageAsRequest(); + assertThat(promptsListRequest.method()).isEqualTo("prompts/list"); + + McpSchema.JSONRPCResponse promptsListResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + promptsListRequest.id(), mockPromptsResult, null); + transport.simulateIncomingMessage(promptsListResponse); + + // Verify the consumer received the expected prompts + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(receivedPrompts).hasSize(1); + assertThat(receivedPrompts.get(0).name()).isEqualTo("test-prompt"); + assertThat(receivedPrompts.get(0).description()).isEqualTo("Test Prompt Description"); + assertThat(receivedPrompts.get(0).arguments()).hasSize(1); + assertThat(receivedPrompts.get(0).arguments().get(0).name()).isEqualTo("arg1"); + }); + + asyncMcpClient.closeGracefully(); + } + }