Skip to content

Commit ad9bc30

Browse files
committed
fix: Support pagination in tools change handler
1 parent 2f94434 commit ad9bc30

File tree

2 files changed

+49
-15
lines changed

2 files changed

+49
-15
lines changed

mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -682,9 +682,17 @@ public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
682682
private NotificationHandler asyncToolsChangeNotificationHandler(
683683
List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers) {
684684
// TODO: params are not used yet
685-
return params -> this.listTools()
686-
.flatMap(listToolsResult -> Flux.fromIterable(toolsChangeConsumers)
687-
.flatMap(consumer -> consumer.apply(listToolsResult.tools()))
685+
return params -> this.listTools().expand(result -> {
686+
if (result.nextCursor() != null) {
687+
return this.listTools(result.nextCursor());
688+
}
689+
return Mono.empty();
690+
}).reduce(new ArrayList<McpSchema.Tool>(), (allTools, result) -> {
691+
allTools.addAll(result.tools());
692+
return allTools;
693+
})
694+
.flatMap(allTools -> Flux.fromIterable(toolsChangeConsumers)
695+
.flatMap(consumer -> consumer.apply(allTools))
688696
.onErrorResume(error -> {
689697
logger.error("Error handling tools list change notification", error);
690698
return Mono.empty();

mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientResponseHandlerTests.java

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import io.modelcontextprotocol.spec.McpSchema;
1818
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
1919
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
20+
import io.modelcontextprotocol.spec.McpSchema.PaginatedRequest;
2021
import io.modelcontextprotocol.spec.McpSchema.Root;
2122
import org.junit.jupiter.api.Test;
2223
import org.junit.jupiter.params.ParameterizedTest;
@@ -109,27 +110,52 @@ void testToolsChangeNotificationHandling() throws JsonProcessingException {
109110

110111
// Create a mock tools list that the server will return
111112
Map<String, Object> inputSchema = Map.of("type", "object", "properties", Map.of(), "required", List.of());
112-
McpSchema.Tool mockTool = new McpSchema.Tool("test-tool", "Test Tool Description",
113+
McpSchema.Tool mockTool = new McpSchema.Tool("test-tool-1", "Test Tool 1 Description",
113114
new ObjectMapper().writeValueAsString(inputSchema));
114-
McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(mockTool), null);
115+
116+
// Create page 1 response with nextPageToken
117+
String nextPageToken = "page2Token";
118+
McpSchema.ListToolsResult mockToolsResult1 = new McpSchema.ListToolsResult(List.of(mockTool), nextPageToken);
115119

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

121-
// Simulate server response to tools/list request
122-
McpSchema.JSONRPCRequest toolsListRequest = transport.getLastSentMessageAsRequest();
123-
assertThat(toolsListRequest.method()).isEqualTo(McpSchema.METHOD_TOOLS_LIST);
125+
// Simulate server response to first tools/list request
126+
McpSchema.JSONRPCRequest toolsListRequest1 = transport.getLastSentMessageAsRequest();
127+
assertThat(toolsListRequest1.method()).isEqualTo(McpSchema.METHOD_TOOLS_LIST);
128+
129+
McpSchema.JSONRPCResponse toolsListResponse1 = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
130+
toolsListRequest1.id(), mockToolsResult1, null);
131+
transport.simulateIncomingMessage(toolsListResponse1);
132+
133+
// Create mock tools for page 2
134+
McpSchema.Tool mockTool2 = new McpSchema.Tool("test-tool-2", "Test Tool 2 Description",
135+
new ObjectMapper().writeValueAsString(inputSchema));
136+
137+
// Create page 2 response with no nextPageToken (last page)
138+
McpSchema.ListToolsResult mockToolsResult2 = new McpSchema.ListToolsResult(List.of(mockTool2), null);
139+
140+
// Simulate server response to second tools/list request with page token
141+
McpSchema.JSONRPCRequest toolsListRequest2 = transport.getLastSentMessageAsRequest();
142+
assertThat(toolsListRequest2.method()).isEqualTo(McpSchema.METHOD_TOOLS_LIST);
143+
144+
// Verify the page token was included in the request
145+
PaginatedRequest params = (PaginatedRequest) toolsListRequest2.params();
146+
assertThat(params).isNotNull();
147+
assertThat(params.cursor()).isEqualTo(nextPageToken);
124148

125-
McpSchema.JSONRPCResponse toolsListResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
126-
toolsListRequest.id(), mockToolsResult, null);
127-
transport.simulateIncomingMessage(toolsListResponse);
149+
McpSchema.JSONRPCResponse toolsListResponse2 = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
150+
toolsListRequest2.id(), mockToolsResult2, null);
151+
transport.simulateIncomingMessage(toolsListResponse2);
128152

129-
// Verify the consumer received the expected tools
130-
assertThat(receivedTools).hasSize(1);
131-
assertThat(receivedTools.get(0).name()).isEqualTo("test-tool");
132-
assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool Description");
153+
// Verify the consumer received all expected tools from both pages
154+
assertThat(receivedTools).hasSize(2);
155+
assertThat(receivedTools.get(0).name()).isEqualTo("test-tool-1");
156+
assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool 1 Description");
157+
assertThat(receivedTools.get(1).name()).isEqualTo("test-tool-2");
158+
assertThat(receivedTools.get(1).description()).isEqualTo("Test Tool 2 Description");
133159

134160
asyncMcpClient.closeGracefully();
135161
}

0 commit comments

Comments
 (0)