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/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java index a4bacde7f4..a11025a72c 100644 --- a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java +++ b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java @@ -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() { @@ -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 @@ -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(); @@ -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)); @@ -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)); 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..24798d051c 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..348719dd37 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 @@ -16,18 +16,20 @@ import java.math.BigInteger; import java.util.ArrayList; import java.util.Base64; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; 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; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SchemaIndex; @@ -87,19 +89,40 @@ public final class McpService { private static final TraitKey ONE_OF_TRAIT = TraitKey.get(OneOfTrait.class); - private final Map tools; - private final Map prompts; + // The tool, prompt, proxy, and service registries are held as immutable snapshots behind volatile + // references (copy-on-write). Readers (tools/list, prompts/list, tools/call dispatch, shutdown) + // read the current snapshot with no locking and always see a complete, consistent map. Every + // mutation (proxy init, dynamic add, and tools/list_changed refresh) rebuilds the affected + // snapshot under registryLock and publishes it atomically, so concurrent mutators cannot lose + // each other's updates and readers never observe a half-updated registry. Network I/O + // (proxy.listTools()/listPrompts()) is always performed outside the lock. + private final Object registryLock = new Object(); + private volatile Map tools; + private volatile Map prompts; + private volatile Map proxies; + private volatile Map services; private final String serviceName; private final String version; - private final Map proxies; - private final Map services; private final AtomicReference initializeRequest = new AtomicReference<>(); private final ToolFilter toolFilter; private final AtomicReference proxiesInitialized = new AtomicReference<>(false); private final McpMetricsObserver metricsObserver; private final SchemaIndex schemaIndex; private final McpServerInterceptor interceptor; - private Consumer notificationWriter; + // Set once via setNotificationWriter() and read later from the refresh/proxy paths on other + // threads, so it is volatile for safe publication. + private volatile 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. + // A single thread also serializes refreshes from different proxies with each other. It is a daemon + // thread that lives for the process; McpService has no explicit lifecycle so it is never shut down. + private final ExecutorService toolRefreshExecutor = + Executors.newSingleThreadExecutor(r -> { + var t = new Thread(r, "mcp-tools-refresh"); + t.setDaemon(true); + return t; + }); McpService( Map services, @@ -110,14 +133,21 @@ public final class McpService { McpMetricsObserver metricsObserver, McpServerInterceptor interceptor ) { - this.services = services; + // Only services needs copying: it is supplied by the builder, which may still hold or mutate + // it. The tools, prompts, and proxies maps are all built fresh here and never referenced + // again, so snapshot() can wrap them without copying. + this.services = snapshot(new LinkedHashMap<>(services)); this.schemaIndex = SchemaIndex.compose(services.values().stream().map(Service::schemaIndex).toArray(SchemaIndex[]::new)); - this.tools = createTools(services); - this.prompts = new ConcurrentHashMap<>(PromptLoader.loadPrompts(services.values())); + this.tools = snapshot(createTools(services)); + this.prompts = snapshot(PromptLoader.loadPrompts(services.values())); this.serviceName = name; this.version = version; - this.proxies = proxyList.stream().collect(Collectors.toMap(McpServerProxy::name, p -> p)); + var proxyMap = new LinkedHashMap(); + for (var proxy : proxyList) { + proxyMap.put(proxy.name(), proxy); + } + this.proxies = snapshot(proxyMap); this.toolFilter = toolFilter; this.metricsObserver = metricsObserver; this.interceptor = interceptor; @@ -328,8 +358,13 @@ private JsonRpcResponse handlePing(JsonRpcRequest req) { } private JsonRpcResponse handlePromptsList(JsonRpcRequest req) { + var promptValues = prompts.values(); + var promptInfos = new ArrayList(promptValues.size()); + for (var prompt : promptValues) { + promptInfos.add(prompt.promptInfo()); + } var result = ListPromptsResult.builder() - .prompts(prompts.values().stream().map(Prompt::promptInfo).toList()) + .prompts(promptInfos) .build(); return createSuccessResponse(req.getId(), result); } @@ -349,12 +384,15 @@ private JsonRpcResponse handlePromptsGet(JsonRpcRequest req) { } private JsonRpcResponse handleToolsList(JsonRpcRequest req, ProtocolVersion protocolVersion) { + var toolValues = tools.values(); + var toolInfos = new ArrayList(toolValues.size()); + for (var tool : toolValues) { + if (toolFilter.allowTool(tool.serverId(), tool.toolInfo().getName())) { + toolInfos.add(extractToolInfo(tool, protocolVersion)); + } + } var result = ListToolsResult.builder() - .tools(tools.values() - .stream() - .filter(t -> toolFilter.allowTool(t.serverId(), t.toolInfo().getName())) - .map(tool -> extractToolInfo(tool, protocolVersion)) - .toList()) + .tools(toolInfos) .build(); return createSuccessResponse(req.getId(), result); } @@ -552,16 +590,12 @@ private Consumer createProxyNotificationWriter( Consumer baseNotificationWriter ) { return notification -> { - // 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, so doing it inline would deadlock the reader. + toolRefreshExecutor.execute(() -> refreshProxyTools(proxy)); } // Forward the notification if (baseNotificationWriter != null) { @@ -570,6 +604,65 @@ 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). The network fetch happens outside + * {@code registryLock}; only the in-memory snapshot swap is locked. 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) { + List proxyTools; + try { + proxyTools = proxy.listTools(); + } catch (Exception e) { + LOG.error("Failed to re-fetch tools from proxy: {}", proxy.name(), e); + return; + } + // Fast path: the proxy reports no tools and has none currently registered, so there is + // nothing to add and nothing to prune. Reads the current snapshot lock-free and avoids the + // set allocation, map copy, and lock entirely. (If a tool for this proxy is added + // concurrently right after this check, that add publishes it and a later refresh reconciles.) + if (proxyTools.isEmpty() && !hasToolsFor(proxy)) { + return; + } + synchronized (registryLock) { + Set newNames = new HashSet<>(); + // LinkedHashMap so a refresh keeps the existing order of every other tool: re-put entries + // stay in place and genuinely new tools append, rather than the whole listing reshuffling + // to hash order on each tools/list_changed. + var next = new LinkedHashMap<>(tools); + for (var toolInfo : proxyTools) { + newNames.add(toolInfo.getName()); + next.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy)); + } + next.entrySet() + .removeIf(entry -> entry.getValue().proxy() == proxy && !newNames.contains(entry.getKey())); + tools = snapshot(next); + } + } + + /** Whether any currently registered tool belongs to the given proxy. Reads the current snapshot. */ + private boolean hasToolsFor(McpServerProxy proxy) { + for (var tool : tools.values()) { + if (tool.proxy() == proxy) { + return true; + } + } + return false; + } + + /** + * Publishes a registry snapshot by wrapping the given map unmodifiable. The caller hands off + * ownership: the argument must be a freshly built map that is never mutated or retained after + * this call, since it becomes the live snapshot without being copied. Callers build these as + * {@link LinkedHashMap}s so {@code tools/list} and {@code prompts/list} return a stable, + * deterministic insertion order across refreshes and dynamic additions. + */ + private static Map snapshot(Map ownedMap) { + return Collections.unmodifiableMap(ownedMap); + } + /** * Starts proxies without initializing them. */ @@ -605,24 +698,52 @@ public void initializeProxies(Consumer responseWriter) { var proxyNotificationWriter = createProxyNotificationWriter(proxy, notificationWriter); proxy.initialize(responseWriter, proxyNotificationWriter, initRequest, protocolVersion); } + // Isolate each proxy: a failure fetching one proxy's tools or prompts must not abort + // discovery for the rest. + registerProxyListing(proxy); + } + } + } + + /** + * Fetches a proxy's tools and prompts (outside {@code registryLock}) and merges them into the + * registries under the lock. A failure fetching either list is logged and skipped so it cannot + * abort discovery for other proxies. + */ + private void registerProxyListing(McpServerProxy proxy) { + List proxyTools = null; + try { + proxyTools = proxy.listTools(); + } catch (Exception e) { + LOG.error("Failed to fetch tools from proxy: {}", proxy.name(), e); + } + + List proxyPrompts = null; + try { + proxyPrompts = proxy.listPrompts(); + } catch (Exception e) { + LOG.error("Failed to fetch prompts from proxy: {}", proxy.name(), e); + } + + if (proxyTools == null && proxyPrompts == null) { + return; + } - List proxyTools = proxy.listTools(); + synchronized (registryLock) { + if (proxyTools != null) { + var nextTools = new LinkedHashMap<>(tools); for (var toolInfo : proxyTools) { - tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy)); + nextTools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy)); } - - // Fetch and register prompts from proxy - try { - List proxyPrompts = proxy.listPrompts(); - for (var promptInfo : proxyPrompts) { - var normalizedName = PromptLoader.normalize(promptInfo.getName()); - if (!prompts.containsKey(normalizedName)) { - prompts.put(normalizedName, new Prompt(promptInfo, proxy)); - } - } - } catch (Exception e) { - LOG.error("Failed to fetch prompts from proxy: " + proxy.name(), e); + tools = snapshot(nextTools); + } + if (proxyPrompts != null) { + var nextPrompts = new LinkedHashMap<>(prompts); + for (var promptInfo : proxyPrompts) { + var normalizedName = PromptLoader.normalize(promptInfo.getName()); + nextPrompts.putIfAbsent(normalizedName, new Prompt(promptInfo, proxy)); } + prompts = snapshot(nextPrompts); } } } @@ -638,39 +759,32 @@ public JsonRpcRequest getInitializeRequest() { * Adds a new service and updates the tools map. */ public void addNewService(String id, Service service) { - services.put(id, service); - tools.putAll(createTools(Map.of(id, service))); + var newTools = createTools(Map.of(id, service)); + synchronized (registryLock) { + var nextServices = new LinkedHashMap<>(services); + nextServices.put(id, service); + services = snapshot(nextServices); + + var nextTools = new LinkedHashMap<>(tools); + nextTools.putAll(newTools); + tools = snapshot(nextTools); + } } public void addNewProxy( McpServerProxy mcpServerProxy, Consumer responseWriter ) { - proxies.put(mcpServerProxy.name(), mcpServerProxy); + synchronized (registryLock) { + var nextProxies = new LinkedHashMap<>(proxies); + nextProxies.put(mcpServerProxy.name(), mcpServerProxy); + proxies = snapshot(nextProxies); + } mcpServerProxy.start(); - try { - List proxyTools = mcpServerProxy.listTools(); - for (var toolInfo : proxyTools) { - tools.put(toolInfo.getName(), new Tool(toolInfo, mcpServerProxy.name(), mcpServerProxy)); - } - } catch (Exception e) { - LOG.error("Failed to fetch tools from proxy", e); - } - - // Also fetch prompts from the new proxy - try { - List proxyPrompts = mcpServerProxy.listPrompts(); - for (var promptInfo : proxyPrompts) { - var normalizedName = PromptLoader.normalize(promptInfo.getName()); - if (!prompts.containsKey(normalizedName)) { - prompts.put(normalizedName, new Prompt(promptInfo, mcpServerProxy)); - } - } - } catch (Exception e) { - LOG.error("Failed to fetch prompts from proxy: " + mcpServerProxy.name(), e); - } + // Fetches tools/prompts (network I/O) outside the lock, then swaps under it. + registerProxyListing(mcpServerProxy); } /** @@ -681,7 +795,8 @@ public boolean containsMcpServer(String id) { } /** - * Returns all registered proxies. + * Returns an immutable snapshot of the registered proxies at the time of the call. Subsequent + * additions via {@link #addNewProxy} are not reflected in a previously returned snapshot. */ public Map getProxies() { return proxies; @@ -799,7 +914,7 @@ private JsonRpcResponse createErrorResponse(JsonRpcRequest req, String s) { } private Map createTools(Map services) { - var tools = new ConcurrentHashMap(); + var tools = new LinkedHashMap(); for (var entry : services.entrySet()) { var id = entry.getKey(); var service = entry.getValue(); 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..e6d130a6c1 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,17 @@ 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. orTimeout() + // completes responseFuture itself on timeout, so the caller (which holds responseFuture) sees + // the TimeoutException; the derived stage exists only to remove 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..c5c69f298f 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..cc31127886 --- /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 to 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/McpServerTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java index 326fef7d86..94efc6e1ec 100644 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java @@ -5,6 +5,7 @@ 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.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -1601,7 +1602,7 @@ void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { } @Test - void testToolsListChangedNotificationInvalidatesCache() { + void testToolsListChangedNotificationInvalidatesCache() throws InterruptedException { var callCounter = new AtomicInteger(0); var mockProxy = new CacheTestProxy(callCounter); @@ -1626,7 +1627,7 @@ void testToolsListChangedNotificationInvalidatesCache() { assertTrue(mockProxy.getSentNotifications().contains("notifications/initialized"), "notifications/initialized should be sent during initialization"); - // First tools/list - fetches from proxy + // First tools/list - reads the tool set fetched from the proxy during initialize. var toolsRequest = JsonRpcRequest.builder() .method("tools/list") .id(Document.of(2)) @@ -1634,13 +1635,15 @@ void testToolsListChangedNotificationInvalidatesCache() { .jsonrpc("2.0") .build(); service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(1, callCounter.get(), "First call should fetch from proxy"); + assertEquals(1, callCounter.get(), "Only initialize should have fetched from the proxy so far"); - // Second tools/list - uses cache + // Second tools/list - still just reads the registry, no proxy fetch. service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(1, callCounter.get(), "Second call should use cache"); + assertEquals(1, callCounter.get(), "tools/list must not fetch from the proxy on its own"); - // Send tools/list_changed notification + // Send tools/list_changed notification. The refresh runs asynchronously off the notifying + // thread (so the transport reader thread can't deadlock), so it fetches from the proxy a + // little later rather than inline. var notification = JsonRpcRequest.builder() .method("notifications/tools/list_changed") .params(Document.of(Map.of())) @@ -1652,13 +1655,17 @@ void testToolsListChangedNotificationInvalidatesCache() { assertEquals(1, notifications.size()); assertEquals("notifications/tools/list_changed", notifications.get(0).getMethod()); - // Third tools/list - should refresh from proxy - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(2, callCounter.get(), "Third call should refresh after notification"); + // The async refresh must eventually fetch from the proxy exactly once. + long deadline = System.nanoTime() + SECONDS.toNanos(5); + while (callCounter.get() < 2 && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertEquals(2, callCounter.get(), "list_changed should trigger exactly one refresh fetch"); - // Fourth tools/list - uses cache again (counter should NOT increment) + // Further tools/list calls just read the refreshed registry, no additional proxy fetch. + service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(2, callCounter.get(), "Fourth call should use cache (not increment to 3)"); + assertEquals(2, callCounter.get(), "tools/list must not fetch from the proxy after the refresh"); } @Test 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..69006b214d --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java @@ -0,0 +1,243 @@ +/* + * 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.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; +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); + private final String name; + + FakeProxy(String name, List initial) { + this.name = name; + 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 name; + } + + 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("fake", 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("fake", 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)); + + // The snapshot swap happens after listTools() returns, so poll for the expected state. + assertEventually(() -> List.of("new1", "old1").equals(listToolNames(service)), + "expected [new1, old1] but was " + listToolNames(service)); + } + + @Test + void concurrentRefreshAddProxyAndListDoNotLoseUpdatesOrThrow() throws Exception { + // Hammer the registry from multiple threads: repeated tools/list_changed refreshes on an + // existing proxy, dynamic addNewProxy calls, and concurrent tools/list reads. With + // copy-on-write under a single lock, reads must never throw and every added proxy's tool + // must be present at the end (no lost updates). + var refreshProxy = new FakeProxy("refresher", List.of(tool("r0"))); + var service = initializedService(refreshProxy); + + int adders = 4; + int refreshers = 4; + int readers = 4; + int iterations = 200; + var barrier = new CyclicBarrier(adders + refreshers + readers); + var error = new AtomicReference(); + var threads = new ArrayList(); + + for (int a = 0; a < adders; a++) { + final int id = a; + threads.add(new Thread(() -> { + try { + barrier.await(); + for (int i = 0; i < iterations; i++) { + var p = new FakeProxy("added-" + id + "-" + i, List.of(tool("added-" + id + "-" + i))); + service.addNewProxy(p, r -> {}); + } + } catch (Throwable t) { + error.compareAndSet(null, t); + } + })); + } + for (int r = 0; r < refreshers; r++) { + threads.add(new Thread(() -> { + try { + barrier.await(); + for (int i = 0; i < iterations; i++) { + refreshProxy.toolSet = List.of(tool("r" + i)); + service.refreshProxyTools(refreshProxy); + } + } catch (Throwable t) { + error.compareAndSet(null, t); + } + })); + } + for (int r = 0; r < readers; r++) { + threads.add(new Thread(() -> { + try { + barrier.await(); + for (int i = 0; i < iterations; i++) { + listToolNames(service); + } + } catch (Throwable t) { + error.compareAndSet(null, t); + } + })); + } + + threads.forEach(Thread::start); + for (var t : threads) { + t.join(30_000); + } + + if (error.get() != null) { + throw new AssertionError("concurrent access threw", error.get()); + } + + // Every proxy added by every adder thread must have its tool registered (no lost updates). + var finalNames = listToolNames(service); + for (int id = 0; id < adders; id++) { + var expected = "added-" + id + "-" + (iterations - 1); + assertTrue(finalNames.contains(expected), + "missing tool from a concurrently added proxy: " + expected); + } + } + + private static void assertEventually(BooleanSupplier condition, String message) + throws InterruptedException { + long deadline = System.nanoTime() + SECONDS.toNanos(5); + while (System.nanoTime() < deadline) { + if (condition.getAsBoolean()) { + return; + } + Thread.sleep(10); + } + throw new AssertionError(message); + } +} 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(); + } + } +}