Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
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 @@ -5,6 +5,7 @@
package io.modelcontextprotocol.server;

import java.time.Duration;
import java.util.Base64;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
Expand Down Expand Up @@ -122,6 +123,11 @@ public class McpAsyncServer {

private McpUriTemplateManagerFactory uriTemplateManagerFactory = new DefaultMcpUriTemplateManagerFactory();

private final TypeRef<McpSchema.PaginatedRequest> PAGINATED_REQUEST_TYPE_REF = new TypeRef<>() {
};

private static final int PAGE_SIZE = 10;

/**
* Create a new McpAsyncServer with the given transport provider and capabilities.
* @param mcpTransportProvider The transport layer implementation for MCP
Expand Down Expand Up @@ -537,9 +543,26 @@ public Mono<Void> notifyToolsListChanged() {

private McpRequestHandler<McpSchema.ListToolsResult> toolsListRequestHandler() {
return (exchange, params) -> {
List<Tool> tools = this.tools.stream().map(McpServerFeatures.AsyncToolSpecification::tool).toList();
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
var cursor = paginatedRequest != null ? paginatedRequest.cursor() : null;

var mapSize = this.tools.size();
var mapHash = this.tools.hashCode();

return handleCursor(cursor, mapSize, mapHash).map(requestedStartIndex -> {
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);

return Mono.just(McpSchema.ListToolsResult.builder(tools).build());
var nextCursor = getCursor(endIndex, mapSize, mapHash);

var resultList = this.tools.stream()
.skip(startIndex)
.limit(endIndex - startIndex)
.map(McpServerFeatures.AsyncToolSpecification::tool)
.toList();

return McpSchema.ListToolsResult.builder(resultList).nextCursor(nextCursor).build();
});
};
}

Expand Down Expand Up @@ -787,21 +810,53 @@ private McpRequestHandler<Object> resourcesUnsubscribeRequestHandler() {

private McpRequestHandler<McpSchema.ListResourcesResult> resourcesListRequestHandler() {
return (exchange, params) -> {
var resourceList = this.resources.values()
.stream()
.map(McpServerFeatures.AsyncResourceSpecification::resource)
.toList();
return Mono.just(McpSchema.ListResourcesResult.builder(resourceList).build());
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
var cursor = paginatedRequest != null ? paginatedRequest.cursor() : null;

var mapSize = this.resources.size();
var mapHash = this.resources.hashCode();

return handleCursor(cursor, mapSize, mapHash).map(requestedStartIndex -> {
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);

var nextCursor = getCursor(endIndex, mapSize, mapHash);

var resultList = this.resources.values()
.stream()
.skip(startIndex)
.limit(endIndex - startIndex)
.map(McpServerFeatures.AsyncResourceSpecification::resource)
.toList();

return McpSchema.ListResourcesResult.builder(resultList).nextCursor(nextCursor).build();
});
};
}

private McpRequestHandler<McpSchema.ListResourceTemplatesResult> resourceTemplateListRequestHandler() {
return (exchange, params) -> {
var resourceList = this.resourceTemplates.values()
.stream()
.map(McpServerFeatures.AsyncResourceTemplateSpecification::resourceTemplate)
.toList();
return Mono.just(McpSchema.ListResourceTemplatesResult.builder(resourceList).build());
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
var cursor = paginatedRequest != null ? paginatedRequest.cursor() : null;

var mapSize = this.resourceTemplates.size();
var mapHash = this.resourceTemplates.hashCode();

return handleCursor(cursor, mapSize, mapHash).map(requestedStartIndex -> {
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);

var nextCursor = getCursor(endIndex, mapSize, mapHash);

var resultList = this.resourceTemplates.values()
.stream()
.skip(startIndex)
.limit(endIndex - startIndex)
.map(McpServerFeatures.AsyncResourceTemplateSpecification::resourceTemplate)
.toList();

return McpSchema.ListResourceTemplatesResult.builder(resultList).nextCursor(nextCursor).build();
});
};
}

Expand Down Expand Up @@ -942,17 +997,27 @@ public Mono<Void> sendElicitationComplete(String sessionId,

private McpRequestHandler<McpSchema.ListPromptsResult> promptsListRequestHandler() {
return (exchange, params) -> {
// TODO: Implement pagination
// McpSchema.PaginatedRequest request = objectMapper.convertValue(params,
// new TypeReference<McpSchema.PaginatedRequest>() {
// });
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
var cursor = paginatedRequest != null ? paginatedRequest.cursor() : null;

var promptList = this.prompts.values()
.stream()
.map(McpServerFeatures.AsyncPromptSpecification::prompt)
.toList();
var mapSize = this.prompts.size();
var mapHash = this.prompts.hashCode();

return Mono.just(McpSchema.ListPromptsResult.builder(promptList).build());
return handleCursor(cursor, mapSize, mapHash).map(requestedStartIndex -> {
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);

var nextCursor = getCursor(endIndex, mapSize, mapHash);

var resultList = this.prompts.values()
.stream()
.skip(startIndex)
.limit(endIndex - startIndex)
.map(McpServerFeatures.AsyncPromptSpecification::prompt)
.toList();

return McpSchema.ListPromptsResult.builder(resultList).nextCursor(nextCursor).build();
});
};
}

Expand Down Expand Up @@ -1106,4 +1171,84 @@ void setProtocolVersions(List<String> protocolVersions) {
this.protocolVersions = protocolVersions;
}

// ---------------------------------------
// Cursor Handling for paginated requests
// ---------------------------------------

/**
* Handles the cursor by decoding, validating and reading the index of it.
* @param cursor the base64 representation of the cursor.
* @param mapSize the size of the map from which the values should be read.
* @param mapHash the hash of the map to compare the cursor value to.
* @return a {@link Mono} which contains the index to which the cursor points.
*/
private Mono<Integer> handleCursor(String cursor, int mapSize, int mapHash) {
if (cursor == null) {
return Mono.just(0);
}

var decodedCursor = decodeCursor(cursor);

if (!isCursorValid(decodedCursor, mapSize, mapHash)) {
return Mono.error(McpError.builder(ErrorCodes.INVALID_PARAMS).message("Invalid cursor").build());
}

return Mono.just(getCursorIndex(decodedCursor));
}

private String getCursor(int endIndex, int mapSize, int mapHash) {
if (endIndex >= mapSize) {
return null;
}
return encodeCursor(endIndex, mapHash);
}

private int getCursorIndex(String cursor) {
return Integer.parseInt(cursor.split(":")[0]);
}

private boolean isCursorValid(String cursor, int maxPageSize, int currentHash) {
var cursorElements = cursor.split(":");

if (cursorElements.length != 2) {
logger.debug("Length of elements in cursor doesn't match expected number. Cursor: {} Actual number: {}",
cursor, cursorElements.length);
return false;
}

int index;
int hash;

try {
index = Integer.parseInt(cursorElements[0]);
hash = Integer.parseInt(cursorElements[1]);
}
catch (NumberFormatException e) {
logger.debug("Failed to parse cursor elements.");
return false;
}

if (index < 0 || index > maxPageSize) {
logger.debug("Cursor boundaries are invalid.");
return false;
}

if (hash != currentHash) {
logger.debug("Cursor not valid, anymore.");
return false;
}

return true;
}

private String encodeCursor(int index, int hash) {
var cursor = index + ":" + hash;

return Base64.getEncoder().encodeToString(cursor.getBytes());
}

private String decodeCursor(String base64Cursor) {
return new String(Base64.getDecoder().decode(base64Cursor));
}

}
Loading
Loading