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
This commit is contained in:
Christian Tzolov 2024-12-19 13:51:30 +01:00 • committed by Christian Tzolov
commit 383d6a7eff
6 changed files with 165 additions and 26 deletions

View file

@ -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<Consumer<List<McpSchema.Tool>>> 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<Consumer<List<McpSchema.Tool>>> 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:

View file

@ -98,7 +98,14 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.awaitility</groupId>
<artifactId>awaitility</artifactId>
<version>4.2.0</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>
</project>

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.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<Supplier<List<Root>>> rootsListProviders, boolean rootsListChangedNotification) {
List<Supplier<List<Root>>> rootsListProviders, boolean rootsListChangedNotification,
List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers) {
Map<String, RequestHandler> 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<String, NotificationHandler> 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<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers) {
return new NotificationHandler() {
@Override
public Mono<Void> handle(Object params) {
// TODO: add support for cursor/pagination
return listTools().flatMap(listToolsResult -> Mono.fromRunnable(() -> {
for (Consumer<List<McpSchema.Tool>> 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<Void>
}
};
};
/**
* The initialization phase MUST be the first interaction between client and server.
* During this phase, the client and server:

View file

@ -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<Supplier<List<Root>>> rootsListProviders = new ArrayList<>();
private List<Consumer<List<McpSchema.Tool>>> 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<List<McpSchema.Tool>> 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);
}
}

View file

@ -183,7 +183,7 @@ public abstract class AbstractMcpAsyncClientTests {
var transport = createMcpTransport();
List<Supplier<List<Root>>> 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();

View file

@ -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<List<Root>> rootsListProvider = () -> List.of(new Root("file:///test/path", "test-root"));
// Create a list to store received tools for verification
List<McpSchema.Tool> receivedTools = new ArrayList<>();
asyncMcpClient = McpClient.using(transport).rootsListProvider(rootsListProvider).async();
}
// Create a consumer that will be called when tools change
Consumer<List<McpSchema.Tool>> 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<String, Object> 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<List<Root>> 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();
}
}