Skip to content

feat: Add client side handlers to handle pagination + fix: Support pagination in tools change notification handler #306

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -682,9 +682,17 @@ public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
private NotificationHandler asyncToolsChangeNotificationHandler(
List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers) {
// TODO: params are not used yet
return params -> this.listTools()
.flatMap(listToolsResult -> Flux.fromIterable(toolsChangeConsumers)
.flatMap(consumer -> consumer.apply(listToolsResult.tools()))
return params -> this.listTools().expand(result -> {
if (result.nextCursor() != null) {
return this.listTools(result.nextCursor());
}
return Mono.empty();
}).reduce(new ArrayList<McpSchema.Tool>(), (allTools, result) -> {
allTools.addAll(result.tools());
return allTools;
})
.flatMap(allTools -> Flux.fromIterable(toolsChangeConsumers)
.flatMap(consumer -> consumer.apply(allTools))
.onErrorResume(error -> {
logger.error("Error handling tools list change notification", error);
return Mono.empty();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
import io.modelcontextprotocol.spec.McpSchema.PaginatedRequest;
import io.modelcontextprotocol.spec.McpSchema.Root;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
Expand Down Expand Up @@ -109,27 +110,52 @@ void testToolsChangeNotificationHandling() throws JsonProcessingException {

// 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",
McpSchema.Tool mockTool = new McpSchema.Tool("test-tool-1", "Test Tool 1 Description",
new ObjectMapper().writeValueAsString(inputSchema));
McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(mockTool), null);

// Create page 1 response with nextPageToken
String nextPageToken = "page2Token";
McpSchema.ListToolsResult mockToolsResult1 = new McpSchema.ListToolsResult(List.of(mockTool), nextPageToken);

// Simulate server sending tools/list_changed notification
McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED, null);
transport.simulateIncomingMessage(notification);

// Simulate server response to tools/list request
McpSchema.JSONRPCRequest toolsListRequest = transport.getLastSentMessageAsRequest();
assertThat(toolsListRequest.method()).isEqualTo(McpSchema.METHOD_TOOLS_LIST);
// Simulate server response to first tools/list request
McpSchema.JSONRPCRequest toolsListRequest1 = transport.getLastSentMessageAsRequest();
assertThat(toolsListRequest1.method()).isEqualTo(McpSchema.METHOD_TOOLS_LIST);

McpSchema.JSONRPCResponse toolsListResponse1 = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
toolsListRequest1.id(), mockToolsResult1, null);
transport.simulateIncomingMessage(toolsListResponse1);

// Create mock tools for page 2
McpSchema.Tool mockTool2 = new McpSchema.Tool("test-tool-2", "Test Tool 2 Description",
new ObjectMapper().writeValueAsString(inputSchema));

// Create page 2 response with no nextPageToken (last page)
McpSchema.ListToolsResult mockToolsResult2 = new McpSchema.ListToolsResult(List.of(mockTool2), null);

// Simulate server response to second tools/list request with page token
McpSchema.JSONRPCRequest toolsListRequest2 = transport.getLastSentMessageAsRequest();
assertThat(toolsListRequest2.method()).isEqualTo(McpSchema.METHOD_TOOLS_LIST);

// Verify the page token was included in the request
PaginatedRequest params = (PaginatedRequest) toolsListRequest2.params();
assertThat(params).isNotNull();
assertThat(params.cursor()).isEqualTo(nextPageToken);

McpSchema.JSONRPCResponse toolsListResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
toolsListRequest.id(), mockToolsResult, null);
transport.simulateIncomingMessage(toolsListResponse);
McpSchema.JSONRPCResponse toolsListResponse2 = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
toolsListRequest2.id(), mockToolsResult2, null);
transport.simulateIncomingMessage(toolsListResponse2);

// Verify the consumer received the expected tools
assertThat(receivedTools).hasSize(1);
assertThat(receivedTools.get(0).name()).isEqualTo("test-tool");
assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool Description");
// Verify the consumer received all expected tools from both pages
assertThat(receivedTools).hasSize(2);
assertThat(receivedTools.get(0).name()).isEqualTo("test-tool-1");
assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool 1 Description");
assertThat(receivedTools.get(1).name()).isEqualTo("test-tool-2");
assertThat(receivedTools.get(1).description()).isEqualTo("Test Tool 2 Description");

asyncMcpClient.closeGracefully();
}
Expand Down