diff --git a/mcp/mcp-schemas/model/main.smithy b/mcp/mcp-schemas/model/main.smithy index d0abb5dfc0..0473261238 100644 --- a/mcp/mcp-schemas/model/main.smithy +++ b/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 { diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java index 22e61867f2..ce5882cd01 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java +++ b/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> notificationConsumer = new AtomicReference<>(); private final AtomicReference> requestNotificationConsumer = new AtomicReference<>(); private final AtomicReference protocolVersion = new AtomicReference<>(ProtocolVersion.defaultVersion()); public List 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 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. + * + *

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 List listPaginated(String method, String errorLabel, PageExtractor extractor) { + List 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 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 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(List items, String nextCursor) {} + + /** Parses a {@code *_/list} result {@code Document} into its items and {@code nextCursor}. */ + @FunctionalInterface + private interface PageExtractor { + Page extract(Document result); } public void initialize( diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java index d22343b2d7..8224d8ba74 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java +++ b/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 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 services, List proxyList, @@ -555,13 +566,10 @@ private Consumer 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 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 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 proxyTools = proxy.listTools(); + Set 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 responseWriter) { proxy.initialize(responseWriter, proxyNotificationWriter, initRequest, protocolVersion); } - List proxyTools = proxy.listTools(); - for (var toolInfo : proxyTools) { - tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy)); + try { + List 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 diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioProxy.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioProxy.java index 94d7388d2f..33b28a0d51 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioProxy.java +++ b/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> 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 arguments; private Map 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 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; } diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java index bf7542ce61..3d79b1bcac 100644 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java @@ -14,6 +14,7 @@ import java.io.OutputStream; import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; +import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; @@ -32,6 +33,8 @@ import software.amazon.smithy.java.json.JsonCodec; import software.amazon.smithy.java.mcp.model.JsonRpcRequest; import software.amazon.smithy.java.mcp.model.JsonRpcResponse; +import software.amazon.smithy.java.mcp.model.ListToolsResult; +import software.amazon.smithy.java.mcp.model.ToolInfo; import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.model.shapes.ShapeType; @@ -574,6 +577,57 @@ public void handle(HttpExchange exchange) throws IOException { } } + @Test + void testListToolsPaginatesAcrossPages() throws IOException { + // Real wire round-trip: the server pages tools/list with a nextCursor, so listTools() must + // follow it and read nextCursor off the actually-deserialized response - not a hand-built + // Document like the McpServerProxy unit tests use. + mockServer.removeContext("/mcp"); + mockServer.createContext("/mcp", exchange -> { + try { + String requestBody = new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8); + JsonRpcRequest request = JsonRpcRequest.builder() + .deserialize(JSON_CODEC.createDeserializer(requestBody.getBytes(StandardCharsets.UTF_8))) + .build(); + var params = request.getParams(); + String cursor = params != null && params.getMember("cursor") != null + ? params.getMember("cursor").asString() + : null; + + var page = ListToolsResult.builder(); + if (cursor == null) { + page.tools(List.of(tool("t1"), tool("t2"))).nextCursor("page2"); + } else { + page.tools(List.of(tool("t3"))); + } + + JsonRpcResponse response = JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .result(Document.of(page.build())) + .build(); + byte[] body = JSON_CODEC.serializeToString(response).getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().set("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + } catch (Exception e) { + exchange.sendResponseHeaders(500, 0); + } finally { + exchange.close(); + } + }); + + List tools = proxy.listTools(); + + assertEquals(List.of("t1", "t2", "t3"), tools.stream().map(ToolInfo::getName).toList()); + } + + private static ToolInfo tool(String name) { + return ToolInfo.builder().name(name).build(); + } + private static class MockMcpHandler implements HttpHandler { @Override public void handle(HttpExchange exchange) throws IOException { diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerProxyTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerProxyTest.java new file mode 100644 index 0000000000..530f3e375e --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerProxyTest.java @@ -0,0 +1,245 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcErrorResponse; +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; + +class McpServerProxyTest { + + /** + * Test proxy that replays a fixed list of canned responses and records every request it received, + * so pagination behaviour (nextCursor -> cursor round-tripping) can be asserted. + */ + private static final class FakeProxy extends McpServerProxy { + private final List requests = new ArrayList<>(); + private final List responses; + private int index = 0; + + FakeProxy(List responses) { + this.responses = responses; + } + + @Override + protected CompletableFuture rpc(JsonRpcRequest request) { + requests.add(request); + return CompletableFuture.completedFuture(responses.get(index++)); + } + + @Override + protected void start() {} + + @Override + protected CompletableFuture shutdown() { + return CompletableFuture.completedFuture(null); + } + + @Override + public String name() { + return "fake"; + } + } + + private static ToolInfo tool(String name) { + return ToolInfo.builder().name(name).build(); + } + + private static PromptInfo prompt(String name) { + return PromptInfo.builder().name(name).build(); + } + + private static JsonRpcResponse toolsResponse(List tools, String nextCursor) { + var result = ListToolsResult.builder().tools(tools); + if (nextCursor != null) { + result.nextCursor(nextCursor); + } + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .result(Document.of(result.build())) + .build(); + } + + private static JsonRpcResponse promptsResponse(List prompts, String nextCursor) { + var result = ListPromptsResult.builder().prompts(prompts); + if (nextCursor != null) { + result.nextCursor(nextCursor); + } + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .result(Document.of(result.build())) + .build(); + } + + private static JsonRpcResponse errorResponse(String message) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .error(JsonRpcErrorResponse.builder().code(-32000).message(message).build()) + .build(); + } + + private static JsonRpcResponse emptyResponse() { + // A malformed response carrying neither a result nor an error. + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .build(); + } + + @Test + void listToolsFollowsNextCursorAcrossPages() { + var proxy = new FakeProxy(List.of( + toolsResponse(List.of(tool("a"), tool("b")), "CURSOR1"), + toolsResponse(List.of(tool("c"), tool("d")), "CURSOR2"), + toolsResponse(List.of(tool("e")), null))); + + var tools = proxy.listTools(); + + assertEquals(List.of("a", "b", "c", "d", "e"), + tools.stream().map(ToolInfo::getName).toList()); + assertEquals(3, proxy.requests.size()); + // First page carries no cursor. + assertNull(proxy.requests.get(0).getParams()); + // Each subsequent page echoes the prior page's nextCursor as the cursor param. + assertEquals("CURSOR1", proxy.requests.get(1).getParams().getMember("cursor").asString()); + assertEquals("CURSOR2", proxy.requests.get(2).getParams().getMember("cursor").asString()); + } + + @Test + void listToolsSinglePageMakesOneCall() { + var proxy = new FakeProxy(List.of( + toolsResponse(List.of(tool("only")), null))); + + var tools = proxy.listTools(); + + assertEquals(1, tools.size()); + assertEquals(1, proxy.requests.size()); + assertNull(proxy.requests.get(0).getParams()); + } + + @Test + void listPromptsFollowsNextCursorAcrossPages() { + var proxy = new FakeProxy(List.of( + promptsResponse(List.of(prompt("p1")), "PC1"), + promptsResponse(List.of(prompt("p2"), prompt("p3")), null))); + + var prompts = proxy.listPrompts(); + + assertEquals(List.of("p1", "p2", "p3"), + prompts.stream().map(PromptInfo::getName).toList()); + assertEquals(2, proxy.requests.size()); + assertEquals("PC1", proxy.requests.get(1).getParams().getMember("cursor").asString()); + } + + @Test + void listToolsAbortsOnRepeatedCursor() { + var proxy = new FakeProxy(List.of( + toolsResponse(List.of(tool("a")), "SAME"), + toolsResponse(List.of(tool("b")), "SAME"))); + + assertThrows(IllegalStateException.class, proxy::listTools); + } + + @Test + void listToolsAbortsAtPageCap() { + // A server that always advances the cursor never trips the repeated-cursor guard, so the + // MAX_LIST_PAGES cap must stop it. Supply 1001 ever-advancing pages; only 1000 are fetched. + var responses = new ArrayList(); + for (int i = 0; i <= 1000; i++) { + responses.add(toolsResponse(List.of(tool("t" + i)), "c" + i)); + } + var proxy = new FakeProxy(responses); + + assertThrows(IllegalStateException.class, proxy::listTools); + assertEquals(1000, proxy.requests.size()); + } + + @Test + void listToolsTreatsBlankCursorAsEndOfList() { + // A server that signals end-of-list with an empty cursor (instead of omitting it) must not + // trigger an extra round-trip or trip the repeated-cursor guard. + var proxy = new FakeProxy(List.of( + toolsResponse(List.of(tool("a"), tool("b")), ""))); + + var tools = proxy.listTools(); + + assertEquals(List.of("a", "b"), tools.stream().map(ToolInfo::getName).toList()); + assertEquals(1, proxy.requests.size()); + } + + @Test + void listToolsAbortsOnCyclingCursor() { + // A -> B -> A is a non-advancing cycle the consecutive-only check would miss; the + // seen-cursor guard must still abort it. + var proxy = new FakeProxy(List.of( + toolsResponse(List.of(tool("a")), "A"), + toolsResponse(List.of(tool("b")), "B"), + toolsResponse(List.of(tool("c")), "A"))); + + assertThrows(IllegalStateException.class, proxy::listTools); + } + + @Test + void listToolsThrowsOnErrorResponse() { + var proxy = new FakeProxy(List.of(errorResponse("boom"))); + + var ex = assertThrows(RuntimeException.class, proxy::listTools); + assertTrue(ex.getMessage().contains("boom")); + } + + @Test + void listToolsThrowsOnErrorOnLaterPage() { + var proxy = new FakeProxy(List.of( + toolsResponse(List.of(tool("a")), "c1"), + errorResponse("kaboom"))); + + assertThrows(RuntimeException.class, proxy::listTools); + assertEquals(2, proxy.requests.size()); + } + + @Test + void listToolsThrowsWhenResponseHasNeitherResultNorError() { + var proxy = new FakeProxy(List.of(emptyResponse())); + + var ex = assertThrows(RuntimeException.class, proxy::listTools); + assertTrue(ex.getMessage().contains("listing tools")); + } + + @Test + void listPromptsAbortsOnRepeatedCursor() { + var proxy = new FakeProxy(List.of( + promptsResponse(List.of(prompt("p1")), "SAME"), + promptsResponse(List.of(prompt("p2")), "SAME"))); + + assertThrows(IllegalStateException.class, proxy::listPrompts); + } + + @Test + void listToolsReturnsImmutableList() { + var proxy = new FakeProxy(List.of(toolsResponse(List.of(tool("a")), null))); + + var tools = proxy.listTools(); + + assertThrows(UnsupportedOperationException.class, () -> tools.add(tool("b"))); + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java new file mode 100644 index 0000000000..29c411fd8c --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java @@ -0,0 +1,149 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.java.mcp.model.JsonRpcResponse; +import software.amazon.smithy.java.mcp.model.ListToolsResult; +import software.amazon.smithy.java.mcp.model.ToolInfo; + +class McpServiceTest { + + /** A proxy whose tool set can change and that records which thread its listTools() ran on. */ + private static final class FakeProxy extends McpServerProxy { + volatile List toolSet; + volatile String lastListToolsThread; + volatile CountDownLatch listToolsLatch = new CountDownLatch(1); + + FakeProxy(List initial) { + this.toolSet = initial; + } + + @Override + public List listTools() { + lastListToolsThread = Thread.currentThread().getName(); + listToolsLatch.countDown(); + return toolSet; + } + + @Override + protected CompletableFuture rpc(JsonRpcRequest request) { + return CompletableFuture.completedFuture(JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId() == null ? Document.of(0) : request.getId()) + .result(Document.of(Map.of())) + .build()); + } + + @Override + protected void start() {} + + @Override + protected CompletableFuture shutdown() { + return CompletableFuture.completedFuture(null); + } + + @Override + public String name() { + return "fake"; + } + + void fireListChanged() { + notify(JsonRpcRequest.builder() + .jsonrpc("2.0") + .method("notifications/tools/list_changed") + .build()); + } + } + + private static ToolInfo tool(String name) { + return ToolInfo.builder().name(name).build(); + } + + /** Builds a service with the fake proxy and drives initialize so its notification writer is wired. */ + private static McpService initializedService(FakeProxy proxy) { + var service = new McpService(Map.of(), + List.of(proxy), + "test", + "1.0", + (s, t) -> true, + null, + McpServerInterceptor.NOOP); + service.handleRequest( + JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method("initialize") + .params(Document.of(Map.of())) + .build(), + r -> {}, + ProtocolVersion.defaultVersion()); + return service; + } + + private static List listToolNames(McpService service) { + var resp = service.handleRequest( + JsonRpcRequest.builder().jsonrpc("2.0").id(Document.of(2)).method("tools/list").build(), + r -> {}, + ProtocolVersion.defaultVersion()); + return resp.getResult() + .asShape(ListToolsResult.builder()) + .getTools() + .stream() + .map(ToolInfo::getName) + .sorted() + .toList(); + } + + @Test + void listChangedRefreshRunsOffTheNotifyingThread() throws Exception { + // Regression for the reader-thread deadlock: a tools/list_changed refresh must NOT run + // listTools() on the thread that delivered the notification (on stdio that is the transport + // reader thread, which must stay free to read the tools/list response). + var proxy = new FakeProxy(List.of(tool("a"))); + initializedService(proxy); + + // initialize() already called listTools() once on this thread; reset for the refresh. + proxy.lastListToolsThread = null; + proxy.listToolsLatch = new CountDownLatch(1); + + proxy.fireListChanged(); + + assertTrue(proxy.listToolsLatch.await(5, SECONDS), "refresh never ran"); + assertNotEquals(Thread.currentThread().getName(), + proxy.lastListToolsThread, + "refresh must not run on the notifying thread"); + assertTrue(proxy.lastListToolsThread != null && proxy.lastListToolsThread.startsWith("mcp-tools-refresh"), + "refresh should run on the dedicated executor thread, was: " + proxy.lastListToolsThread); + } + + @Test + void listChangedRefreshAddsNewToolsAndPrunesStaleOnes() throws Exception { + var proxy = new FakeProxy(List.of(tool("old1"), tool("old2"))); + var service = initializedService(proxy); + assertEquals(List.of("old1", "old2"), listToolNames(service)); + + // Server's set changes: old1 kept, old2 gone, new1 added. + proxy.toolSet = List.of(tool("old1"), tool("new1")); + proxy.listToolsLatch = new CountDownLatch(1); + proxy.fireListChanged(); + assertTrue(proxy.listToolsLatch.await(5, SECONDS)); + Thread.sleep(100); // let the map swap after listTools() returns + + assertEquals(List.of("new1", "old1"), listToolNames(service)); + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioProxyTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioProxyTest.java new file mode 100644 index 0000000000..34367f61de --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioProxyTest.java @@ -0,0 +1,48 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.TimeoutException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledOnOs; +import org.junit.jupiter.api.condition.OS; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; + +class StdioProxyTest { + + @Test + @EnabledOnOs({OS.LINUX, OS.MAC}) + void rpcTimesOutWhenServerStaysSilent() { + // `sleep` accepts the request on stdin but never writes a response, so the request future must + // fail via the per-request timeout rather than blocking the caller forever. + var proxy = StdioProxy.builder() + .name("silent-server") + .command("sleep") + .arguments(List.of("30")) + .timeout(Duration.ofMillis(500)) + .build(); + proxy.start(); + try { + var future = proxy.rpc(JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method("tools/list") + .build()); + + var ex = assertThrows(CompletionException.class, future::join); + assertInstanceOf(TimeoutException.class, ex.getCause()); + } finally { + proxy.shutdown().join(); + } + } +}