Paginate MCP tools/list and prompts/list

mcp/mcp-schemas/model/main.smithy

@@ -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 {
@@ -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 {

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java

@@ -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;
@@ -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;
@@ -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(

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java

@@ -25,6 +25,8 @@
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -101,6 +103,15 @@ public final class McpService {
     private final McpServerInterceptor interceptor;
     private Consumer<JsonRpcRequest> notificationWriter;
 
+    // Runs tools/list_changed refreshes off the transport's reader thread. A synchronous refresh calls
+    // listTools() whose response is read by that same reader thread, so doing it inline deadlocks it.
+    private final ExecutorService toolRefreshExecutor =
+            Executors.newSingleThreadExecutor(r -> {
+                var t = new Thread(r, "mcp-tools-refresh");
+                t.setDaemon(true);
+                return t;
+            });
+
     McpService(
             Map<String, Service> services,
             List<McpServerProxy> proxyList,
@@ -555,13 +566,10 @@ private Consumer<JsonRpcRequest> createProxyNotificationWriter(
             // Check if this is a tools/list_changed notification
             if ("notifications/tools/list_changed".equals(notification.getMethod())) {
                 LOG.debug("Received tools/list_changed notification from proxy: {}", proxy.name());
-                // Remove only this proxy's tools
-                tools.entrySet().removeIf(entry -> entry.getValue().proxy() == proxy);
-                // Re-fetch tools from only this proxy
-                List<ToolInfo> proxyTools = proxy.listTools();
-                for (var toolInfo : proxyTools) {
-                    tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy));
-                }
+                // Refresh on a separate thread. This notification is delivered on the proxy's transport
+                // reader thread, and refreshProxyTools() calls listTools() whose response is read by that
+                // same thread -- doing it inline would deadlock the reader.
+                toolRefreshExecutor.execute(() -> refreshProxyTools(proxy));
             }
             // Forward the notification
             if (baseNotificationWriter != null) {
@@ -570,6 +578,27 @@ private Consumer<JsonRpcRequest> createProxyNotificationWriter(
         };
     }
 
+    /**
+     * Re-fetches a proxy's tools after a {@code tools/list_changed} notification and swaps them into the
+     * registry. Runs off the transport reader thread (see caller). Fetches first so a failed or slow
+     * refresh never wipes the current tools, then adds the new set before pruning this proxy's stale
+     * entries, so a concurrent {@code tools/list} never observes a gap (at worst a brief superset).
+     */
+    void refreshProxyTools(McpServerProxy proxy) {
+        try {
+            List<ToolInfo> proxyTools = proxy.listTools();
+            Set<String> newNames = new HashSet<>();
+            for (var toolInfo : proxyTools) {
+                newNames.add(toolInfo.getName());
+                tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy));
+            }
+            tools.entrySet()
+                    .removeIf(entry -> entry.getValue().proxy() == proxy && !newNames.contains(entry.getKey()));
+        } catch (Exception e) {
+            LOG.error("Failed to re-fetch tools from proxy: " + proxy.name(), e);
+        }
+    }
+
     /**
      * Starts proxies without initializing them.
      */
@@ -606,9 +635,13 @@ public void initializeProxies(Consumer<JsonRpcResponse> responseWriter) {
                     proxy.initialize(responseWriter, proxyNotificationWriter, initRequest, protocolVersion);
                 }
 
-                List<ToolInfo> proxyTools = proxy.listTools();
-                for (var toolInfo : proxyTools) {
-                    tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy));
+                try {
+                    List<ToolInfo> proxyTools = proxy.listTools();
+                    for (var toolInfo : proxyTools) {
+                        tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy));
+                    }
+                } catch (Exception e) {
+                    LOG.error("Failed to fetch tools from proxy: " + proxy.name(), e);
                 }
 
                 // Fetch and register prompts from proxy

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioProxy.java

@@ -5,6 +5,7 @@
 
 package software.amazon.smithy.java.mcp.server;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import java.io.BufferedReader;
@@ -14,6 +15,7 @@
 import java.io.InputStreamReader;
 import java.io.OutputStreamWriter;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -42,6 +44,7 @@ public final class StdioProxy extends McpServerProxy {
     private final Map<String, CompletableFuture<JsonRpcResponse>> pendingRequests = new ConcurrentHashMap<>();
     private volatile boolean running = false;
     private final String name;
+    private final Duration requestTimeout;
 
     private StdioProxy(Builder builder) {
         processBuilder = new ProcessBuilder();
@@ -62,6 +65,7 @@ private StdioProxy(Builder builder) {
         }
 
         this.name = builder.name;
+        this.requestTimeout = builder.timeout != null ? builder.timeout : Duration.ofMinutes(5);
 
         processBuilder.redirectErrorStream(false); // Keep stderr separate
     }
@@ -72,6 +76,7 @@ public static class Builder {
         private List<String> arguments;
         private Map<String, String> environmentVariables;
         private File workingDirectory;
+        private Duration timeout;
 
         public Builder name(String name) {
             this.name = name;
@@ -98,6 +103,16 @@ public Builder workingDirectory(File workingDirectory) {
             return this;
         }
 
+        /**
+         * Per-request timeout: a request that never receives a matching response (e.g. a server that
+         * stays alive but goes silent) fails after this duration instead of blocking the caller
+         * forever. Defaults to 5 minutes, symmetric with {@link HttpMcpProxy}.
+         */
+        public Builder timeout(Duration timeout) {
+            this.timeout = timeout;
+            return this;
+        }
+
         public StdioProxy build() {
             if (command == null || command.isEmpty()) {
                 throw new IllegalArgumentException("Command must be provided");
@@ -158,6 +173,15 @@ public CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
             writeLock.unlock();
         }
 
+        // Fail a request that never receives a matching response (server alive but silent) instead of
+        // blocking the caller forever; symmetric with HttpMcpProxy's request timeout. Also removes the
+        // pending-request entry on any completion (success, error, or timeout). Skipped when the write
+        // above already failed and completed the future.
+        if (!responseFuture.isDone()) {
+            responseFuture.orTimeout(requestTimeout.toMillis(), MILLISECONDS)
+                    .whenComplete((response, error) -> pendingRequests.remove(requestId));
+        }
+
         return responseFuture;
     }