Skip to content

Commit e68bbd8

Browse files
denniskawurekdenniskawurek
committed
feat: Server handles pagination
1 parent 4a95718 commit e68bbd8

3 files changed

Lines changed: 875 additions & 48 deletions

File tree

‎mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java‎

Lines changed: 162 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
package io.modelcontextprotocol.server;
66

77
import java.time.Duration;
8+
import java.util.Base64;
89
import java.util.Collections;
910
import java.util.HashMap;
1011
import java.util.List;
@@ -123,6 +124,11 @@ public class McpAsyncServer {
123124

124125
private McpUriTemplateManagerFactory uriTemplateManagerFactory = new DefaultMcpUriTemplateManagerFactory();
125126

127+
private final TypeRef<McpSchema.PaginatedRequest> PAGINATED_REQUEST_TYPE_REF = new TypeRef<>() {
128+
};
129+
130+
private static final int PAGE_SIZE = 10;
131+
126132
/**
127133
* Create a new McpAsyncServer with the given transport provider and capabilities.
128134
* @param mcpTransportProvider The transport layer implementation for MCP
@@ -537,9 +543,25 @@ public Mono<Void> notifyToolsListChanged() {
537543

538544
private McpRequestHandler<McpSchema.ListToolsResult> toolsListRequestHandler() {
539545
return (exchange, params) -> {
540-
List<Tool> tools = this.tools.stream().map(McpServerFeatures.AsyncToolSpecification::tool).toList();
546+
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
547+
548+
var mapSize = this.tools.size();
549+
var mapHash = this.tools.hashCode();
550+
551+
return handleCursor(paginatedRequest.cursor(), mapSize, mapHash).map(requestedStartIndex -> {
552+
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
553+
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);
541554

542-
return Mono.just(McpSchema.ListToolsResult.builder(tools).build());
555+
var nextCursor = getCursor(endIndex, mapSize, mapHash);
556+
557+
var resultList = this.tools.stream()
558+
.skip(startIndex)
559+
.limit(endIndex - startIndex)
560+
.map(McpServerFeatures.AsyncToolSpecification::tool)
561+
.toList();
562+
563+
return McpSchema.ListToolsResult.builder(resultList).nextCursor(nextCursor).build();
564+
});
543565
};
544566
}
545567

@@ -787,21 +809,51 @@ private McpRequestHandler<Object> resourcesUnsubscribeRequestHandler() {
787809

788810
private McpRequestHandler<McpSchema.ListResourcesResult> resourcesListRequestHandler() {
789811
return (exchange, params) -> {
790-
var resourceList = this.resources.values()
791-
.stream()
792-
.map(McpServerFeatures.AsyncResourceSpecification::resource)
793-
.toList();
794-
return Mono.just(McpSchema.ListResourcesResult.builder(resourceList).build());
812+
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
813+
814+
var mapSize = this.resources.size();
815+
var mapHash = this.resources.hashCode();
816+
817+
return handleCursor(paginatedRequest.cursor(), mapSize, mapHash).map(requestedStartIndex -> {
818+
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
819+
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);
820+
821+
var nextCursor = getCursor(endIndex, mapSize, mapHash);
822+
823+
var resultList = this.resources.values()
824+
.stream()
825+
.skip(startIndex)
826+
.limit(endIndex - startIndex)
827+
.map(McpServerFeatures.AsyncResourceSpecification::resource)
828+
.toList();
829+
830+
return McpSchema.ListResourcesResult.builder(resultList).nextCursor(nextCursor).build();
831+
});
795832
};
796833
}
797834

798835
private McpRequestHandler<McpSchema.ListResourceTemplatesResult> resourceTemplateListRequestHandler() {
799836
return (exchange, params) -> {
800-
var resourceList = this.resourceTemplates.values()
801-
.stream()
802-
.map(McpServerFeatures.AsyncResourceTemplateSpecification::resourceTemplate)
803-
.toList();
804-
return Mono.just(McpSchema.ListResourceTemplatesResult.builder(resourceList).build());
837+
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
838+
839+
var mapSize = this.resourceTemplates.size();
840+
var mapHash = this.resourceTemplates.hashCode();
841+
842+
return handleCursor(paginatedRequest.cursor(), mapSize, mapHash).map(requestedStartIndex -> {
843+
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
844+
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);
845+
846+
var nextCursor = getCursor(endIndex, mapSize, mapHash);
847+
848+
var resultList = this.resourceTemplates.values()
849+
.stream()
850+
.skip(startIndex)
851+
.limit(endIndex - startIndex)
852+
.map(McpServerFeatures.AsyncResourceTemplateSpecification::resourceTemplate)
853+
.toList();
854+
855+
return McpSchema.ListResourceTemplatesResult.builder(resultList).nextCursor(nextCursor).build();
856+
});
805857
};
806858
}
807859

@@ -923,17 +975,26 @@ public Mono<Void> notifyPromptsListChanged() {
923975

924976
private McpRequestHandler<McpSchema.ListPromptsResult> promptsListRequestHandler() {
925977
return (exchange, params) -> {
926-
// TODO: Implement pagination
927-
// McpSchema.PaginatedRequest request = objectMapper.convertValue(params,
928-
// new TypeReference<McpSchema.PaginatedRequest>() {
929-
// });
978+
var paginatedRequest = jsonMapper.convertValue(params, PAGINATED_REQUEST_TYPE_REF);
930979

931-
var promptList = this.prompts.values()
932-
.stream()
933-
.map(McpServerFeatures.AsyncPromptSpecification::prompt)
934-
.toList();
980+
var mapSize = this.prompts.size();
981+
var mapHash = this.prompts.hashCode();
935982

936-
return Mono.just(McpSchema.ListPromptsResult.builder(promptList).build());
983+
return handleCursor(paginatedRequest.cursor(), mapSize, mapHash).map(requestedStartIndex -> {
984+
var startIndex = requestedStartIndex != null ? requestedStartIndex : 0;
985+
var endIndex = Math.min(startIndex + PAGE_SIZE, mapSize);
986+
987+
var nextCursor = getCursor(endIndex, mapSize, mapHash);
988+
989+
var resultList = this.prompts.values()
990+
.stream()
991+
.skip(startIndex)
992+
.limit(endIndex - startIndex)
993+
.map(McpServerFeatures.AsyncPromptSpecification::prompt)
994+
.toList();
995+
996+
return McpSchema.ListPromptsResult.builder(resultList).nextCursor(nextCursor).build();
997+
});
937998
};
938999
}
9391000

@@ -1089,4 +1150,84 @@ void setProtocolVersions(List<String> protocolVersions) {
10891150
this.protocolVersions = protocolVersions;
10901151
}
10911152

1153+
// ---------------------------------------
1154+
// Cursor Handling for paginated requests
1155+
// ---------------------------------------
1156+
1157+
/**
1158+
* Handles the cursor by decoding, validating and reading the index of it.
1159+
* @param cursor the base64 representation of the cursor.
1160+
* @param mapSize the size of the map from which the values should be read.
1161+
* @param mapHash the hash of the map to compare the cursor value to.
1162+
* @return a {@link Mono} which contains the index to which the cursor points.
1163+
*/
1164+
private Mono<Integer> handleCursor(String cursor, int mapSize, int mapHash) {
1165+
if (cursor == null) {
1166+
return Mono.just(0);
1167+
}
1168+
1169+
var decodedCursor = decodeCursor(cursor);
1170+
1171+
if (!isCursorValid(decodedCursor, mapSize, mapHash)) {
1172+
return Mono.error(McpError.builder(ErrorCodes.INVALID_PARAMS).message("Invalid cursor").build());
1173+
}
1174+
1175+
return Mono.just(getCursorIndex(decodedCursor));
1176+
}
1177+
1178+
private String getCursor(int endIndex, int mapSize, int mapHash) {
1179+
if (endIndex >= mapSize) {
1180+
return null;
1181+
}
1182+
return encodeCursor(endIndex, mapHash);
1183+
}
1184+
1185+
private int getCursorIndex(String cursor) {
1186+
return Integer.parseInt(cursor.split(":")[0]);
1187+
}
1188+
1189+
private boolean isCursorValid(String cursor, int maxPageSize, int currentHash) {
1190+
var cursorElements = cursor.split(":");
1191+
1192+
if (cursorElements.length != 2) {
1193+
logger.debug("Length of elements in cursor doesn't match expected number. Cursor: {} Actual number: {}",
1194+
cursor, cursorElements.length);
1195+
return false;
1196+
}
1197+
1198+
int index;
1199+
int hash;
1200+
1201+
try {
1202+
index = Integer.parseInt(cursorElements[0]);
1203+
hash = Integer.parseInt(cursorElements[1]);
1204+
}
1205+
catch (NumberFormatException e) {
1206+
logger.debug("Failed to parse cursor elements.");
1207+
return false;
1208+
}
1209+
1210+
if (index < 0 || index > maxPageSize) {
1211+
logger.debug("Cursor boundaries are invalid.");
1212+
return false;
1213+
}
1214+
1215+
if (hash != currentHash) {
1216+
logger.debug("Cursor not valid, anymore.");
1217+
return false;
1218+
}
1219+
1220+
return true;
1221+
}
1222+
1223+
private String encodeCursor(int index, int hash) {
1224+
var cursor = index + ":" + hash;
1225+
1226+
return Base64.getEncoder().encodeToString(cursor.getBytes());
1227+
}
1228+
1229+
private String decodeCursor(String base64Cursor) {
1230+
return new String(Base64.getDecoder().decode(base64Cursor));
1231+
}
1232+
10921233
}

0 commit comments

Comments
 (0)