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:
parent
9adcc5b9c4
commit
383d6a7eff
6 changed files with 165 additions and 26 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue