Skip to content

Paginate MCP tools/list and prompts/list #1336

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.

Sign up for GitHub

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

Open
sugmanue wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Paginate MCP tools/list and prompts/list #1336

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
896f212
Paginate MCP tools/list and prompts/list
sugmanue Aug 28, 2026
10bac1b
Make McpService tool/prompt/proxy registries thread-safe
sugmanue Aug 28, 2026
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
8 changes: 8 additions & 0 deletions mcp/mcp-schemas/model/main.smithy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,10 @@ structure ServerInfo {

structure ListToolsResult {
tools: ToolInfoList

/// Opaque cursor for the next page of results, per MCP pagination (spec section 5.4).
/// Absent when there are no further pages.
nextCursor: String
}

structure ToolInfo {
Expand Down Expand Up @@ -219,6 +223,10 @@ structure TextContent {

structure ListPromptsResult {
prompts: PromptInfoList

/// Opaque cursor for the next page of results, per MCP pagination (spec section 5.4).
/// Absent when there are no further pages.
nextCursor: String
}

list PromptInfoList {
Expand Down
59 changes: 28 additions & 31 deletions ...p-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,13 +133,18 @@ private record ToolSchemas(Schema inputSchema, Schema outputSchema, JsonNode too
private ToolSchemas getMcpEchoToolSchemas() {
write("tools/list", Document.of(Map.of()));
var responseJson = readRawResponse();
var toolNode = OBJECT_MAPPER.readTree(responseJson).path("result").path("tools").get(0);
var inputSchemaNode = toolNode.path("inputSchema");
var outputSchemaNode = toolNode.path("outputSchema");
return new ToolSchemas(
SCHEMA_FACTORY.getSchema(inputSchemaNode),
SCHEMA_FACTORY.getSchema(outputSchemaNode),
toolNode);
var toolsNode = OBJECT_MAPPER.readTree(responseJson).path("result").path("tools");
for (var toolNode : toolsNode) {
if (toolNode.path("name").asString().equals("McpEcho")) {
var inputSchemaNode = toolNode.path("inputSchema");
var outputSchemaNode = toolNode.path("outputSchema");
return new ToolSchemas(
SCHEMA_FACTORY.getSchema(inputSchemaNode),
SCHEMA_FACTORY.getSchema(outputSchemaNode),
toolNode);
}
}
throw new AssertionError("McpEcho tool not found");
}

private ToolSchemas getCalculateAreaToolSchemas() {
Expand All @@ -163,6 +168,19 @@ private String readRawResponse() {
return assertTimeoutPreemptively(Duration.ofSeconds(5), output::read, "No response within 5 seconds");
}

/** Returns the McpEcho tool's outputSchema node, located by name rather than list position. */
private JsonNode mcpEchoOutputSchemaNode() {
write("tools/list", Document.of(Map.of()));
var responseJson = readRawResponse();
var toolsNode = OBJECT_MAPPER.readTree(responseJson).path("result").path("tools");
for (var toolNode : toolsNode) {
if (toolNode.path("name").asString().equals("McpEcho")) {
return toolNode.path("outputSchema");
}
}
throw new AssertionError("McpEcho tool not found");
}

// ========== Protocol Version Tests ==========

@Test
Expand Down Expand Up @@ -862,14 +880,7 @@ void testDocumentValidatesAgainstSchema(String description, Document documentVal
initializeLatestProtocol();

// Get output schema from raw JSON (using Jackson directly)
write("tools/list", Document.of(Map.of()));
var toolsResponseJson = readRawResponse();
var toolsResponseNode = OBJECT_MAPPER.readTree(toolsResponseJson);
var outputSchemaNode = toolsResponseNode
.path("result")
.path("tools")
.get(0)
.path("outputSchema");
var outputSchemaNode = mcpEchoOutputSchemaNode();

// Create input with the document value
var echoData = new HashMap<String, Document>();
Expand Down Expand Up @@ -1347,14 +1358,7 @@ void testStructuredContentValidatesAgainstOutputSchema() {
initializeLatestProtocol();

// Get output schema from raw JSON (using Jackson directly, not Smithy serializers)
write("tools/list", Document.of(Map.of()));
var toolsResponseJson = readRawResponse();
var toolsResponseNode = OBJECT_MAPPER.readTree(toolsResponseJson);
var outputSchemaNode = toolsResponseNode
.path("result")
.path("tools")
.get(0)
.path("outputSchema");
var outputSchemaNode = mcpEchoOutputSchemaNode();

// Create comprehensive input
var base64Blob = Base64.getEncoder().encodeToString("test".getBytes(StandardCharsets.UTF_8));
Expand Down Expand Up @@ -1396,14 +1400,7 @@ void testAllTypesValidateAgainstOutputSchema() {
initializeLatestProtocol();

// Get output schema from raw JSON (using Jackson directly)
write("tools/list", Document.of(Map.of()));
var toolsResponseJson = readRawResponse();
var toolsResponseNode = OBJECT_MAPPER.readTree(toolsResponseJson);
var outputSchemaNode = toolsResponseNode
.path("result")
.path("tools")
.get(0)
.path("outputSchema");
var outputSchemaNode = mcpEchoOutputSchemaNode();

// Create input with all types
var base64Blob = Base64.getEncoder().encodeToString("binary data".getBytes(StandardCharsets.UTF_8));
Expand Down
128 changes: 97 additions & 31 deletions mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@

package software.amazon.smithy.java.mcp.server;

import static software.amazon.smithy.java.mcp.model.ListPromptsResult.builder;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
Expand All @@ -19,6 +21,7 @@
import software.amazon.smithy.java.logging.InternalLogger;
import software.amazon.smithy.java.mcp.model.JsonRpcRequest;
import software.amazon.smithy.java.mcp.model.JsonRpcResponse;
import software.amazon.smithy.java.mcp.model.ListPromptsResult;
import software.amazon.smithy.java.mcp.model.ListToolsResult;
import software.amazon.smithy.java.mcp.model.PromptInfo;
import software.amazon.smithy.java.mcp.model.ToolInfo;
Expand All @@ -31,46 +34,109 @@ public abstract class McpServerProxy {
private static final InternalLogger LOG = InternalLogger.getLogger(McpServerProxy.class);
private static final AtomicInteger ID_GENERATOR = new AtomicInteger(0);

// Cap list pages so a server that always returns a fresh, advancing cursor fails the call
// instead of looping forever. MCP cursors are opaque and the spec does not guarantee
// termination; at a typical ~30 items/page this bounds a listing at ~30k items.
private static final int MAX_LIST_PAGES = 1000;

private final AtomicReference<Consumer<JsonRpcResponse>> notificationConsumer = new AtomicReference<>();
private final AtomicReference<Consumer<JsonRpcRequest>> requestNotificationConsumer = new AtomicReference<>();
private final AtomicReference<ProtocolVersion> protocolVersion =
new AtomicReference<>(ProtocolVersion.defaultVersion());

public List<ToolInfo> listTools() {
JsonRpcRequest request = JsonRpcRequest.builder()
.method("tools/list")
.id(generateRequestId())
.jsonrpc("2.0")
.build();

return rpc(request).thenApply(response -> {
if (response.getError() != null) {
throw new RuntimeException("Error listing tools: " + response.getError().getMessage());
}
return response.getResult()
.asShape(ListToolsResult.builder())
.getTools()
.stream()
.toList();
}).join();
return listPaginated("tools/list", "listing tools", result -> {
ListToolsResult page = result.asShape(ListToolsResult.builder());
return new Page<>(page.getTools(), page.getNextCursor());
});
}

public List<PromptInfo> listPrompts() {
JsonRpcRequest request = JsonRpcRequest.builder()
.method("prompts/list")
.id(generateRequestId())
.jsonrpc("2.0")
.build();
return rpc(request).thenApply(response -> {
return listPaginated("prompts/list", "listing prompts", result -> {
ListPromptsResult page = result.asShape(ListPromptsResult.builder());
return new Page<>(page.getPrompts(), page.getNextCursor());
});
}

/**
* Maximum number of pages {@link #listTools()} / {@link #listPrompts()} will fetch before
* aborting, a backstop against a server that keeps returning a fresh, advancing cursor and
* never terminates. Subclasses may override to tighten or relax the bound.
*/
protected int maxListPages() {
return MAX_LIST_PAGES;
}

/**
* Drives MCP cursor pagination for a {@code tools/list}-style method: repeatedly calls
* {@code method}, threading the previous page's {@code nextCursor} back as the {@code cursor}
* request param, and accumulates items across all pages in page order until the server stops
* returning a cursor. A single-page server (no {@code nextCursor}) makes exactly one round-trip.
*
* <p>Three guards bound a misbehaving server: an absent or blank {@code nextCursor} ends
* pagination; a previously-seen cursor (including a non-advancing {@code A -> B -> A} cycle)
* aborts; and the page count is capped at {@link #maxListPages()}.
*/
private <T> List<T> listPaginated(String method, String errorLabel, PageExtractor<T> extractor) {
List<T> all = new ArrayList<>();
// Cursors already requested this call, so a repeated or cycling cursor is caught immediately
// rather than only when two identical cursors happen to be adjacent.
Set<String> seenCursors = new HashSet<>();
String cursor = null;
int page = 0;
do {
if (++page > maxListPages()) {
throw new IllegalStateException(
"Aborting " + method + ": server returned more than " + maxListPages()
+ " pages without terminating (possible pagination bug or misbehaving server)");
}

JsonRpcRequest.Builder requestBuilder = JsonRpcRequest.builder()
.method(method)
.id(generateRequestId())
.jsonrpc("2.0");
if (cursor != null) {
requestBuilder.params(Document.of(Map.of("cursor", Document.of(cursor))));
}

JsonRpcResponse response = rpc(requestBuilder.build()).join();
if (response.getError() != null) {
throw new RuntimeException("Error listing prompts: " + response.getError().getMessage());
throw new RuntimeException("Error " + errorLabel + ": " + response.getError().getMessage());
}
return response.getResult()
.asShape(builder())
.getPrompts()
.stream()
.toList();
}).join();

Document result = response.getResult();
if (result == null) {
throw new RuntimeException(
"Error " + errorLabel + ": response contained neither a result nor an error");
}

Page<T> parsed = extractor.extract(result);
all.addAll(parsed.items());

// MCP signals "no more pages" by omitting nextCursor; defensively treat a blank cursor the
// same way, since some servers send "" instead of omitting the field.
String nextCursor = parsed.nextCursor();
if (nextCursor != null && nextCursor.isBlank()) {
nextCursor = null;
}
if (nextCursor != null && !seenCursors.add(nextCursor)) {
throw new IllegalStateException(
"Aborting " + method + ": server repeated a pagination cursor (no forward progress)");
}
cursor = nextCursor;
} while (cursor != null);

LOG.debug("{}: fetched {} item(s) across {} page(s)", method, all.size(), page);
return List.copyOf(all);
}

/** One page of a paginated list: the page's items plus the server's {@code nextCursor} (null when last). */
private record Page<T>(List<T> items, String nextCursor) {}

/** Parses a {@code *_/list} result {@code Document} into its items and {@code nextCursor}. */
@FunctionalInterface
private interface PageExtractor<T> {
Page<T> extract(Document result);
}

public void initialize(
Expand Down
Loading
Loading