diff --git a/examples/mcp-server/README.md b/examples/mcp-server/README.md index 5e32ef0eab..6c46773c7a 100644 --- a/examples/mcp-server/README.md +++ b/examples/mcp-server/README.md @@ -1,5 +1,12 @@ ## Example: MCP Server +This example contains two newline-delimited JSON-RPC servers using the MCP +standard input/output transport: + +- `MCPServerExample` exposes generated Smithy service implementations directly. +- `ProxyMCPExample` starts a Smithy HTTP server on port `8080` and exposes a + `ProxyService` for it through MCP. + ### Usage To use this example as a template, run the following command with @@ -15,34 +22,42 @@ Or smithy init -t mcp-server --url git@github.com:smithy-lang/smithy-java.git ``` -To generate a fat jar which contains all the dependencies required to run -a [Model Context Protocol](https://modelcontextprotocol.io/) ( -MCP) [StdIO](https://modelcontextprotocol.io/docs/concepts/transports#standard-input%2Foutput-stdio) server, -run the following from the root of the project: +The generated server uses the transport-specific `StdioMcpServer` entry point: -```console -gradle build -``` +```java +var mcpServer = StdioMcpServer.builder() + .stdio() + .name("smithy-mcp-server") + .addService("employee-mcp", service) + .build(); -This will generate a fat JAR file at `build/libs/mcp-server-0.0.1-all.jar`. This artifact includes all the necessary -code to create an MCP server that uses the StdIO transport. +mcpServer.start(); +mcpServer.awaitCompletion(); +``` -There are two example implementations included: +To compile both implementations and generate a fat JAR from a Smithy Java +checkout, run: -* `MCPServerExample` : Demonstrates how to build an MCP server by modeling tools as Smithy APIs. +```console +./gradlew :examples:mcp-server:build +``` -* `ProxyMCPExample` : Shows how to create a Proxy MCP Server for any Smithy service. In this example, a Smithy Java - server is started on port 8080, and the MCP server proxies requests to it. +The fat JAR is written to +`examples/mcp-server/build/libs/mcp-server--all.jar`. It contains the +generated service code, both example entry points, and the MCP standard +input/output transport. -You can run the Proxy MCP Server using the following command: +Run the proxy example from the repository root with: -``` -java -cp mcp-server-0.0.1-all.jar software.amazon.smithy.java.example.server.mcp.ProxyMCPExample +```console +java -cp examples/mcp-server/build/libs/mcp-server-*-all.jar \ + software.amazon.smithy.java.example.server.mcp.ProxyMCPExample ``` -To run the direct MCP server example instead, simply replace `ProxyMCPExample` with `MCPServerExample`. +Replace `ProxyMCPExample` with `MCPServerExample` to run the direct service +implementation. -Here's how you might configure the MCP client to invoke the proxy server: +An MCP client can launch the proxy server with a configuration like: ```json { @@ -51,14 +66,10 @@ Here's how you might configure the MCP client to invoke the proxy server: "command": "java", "args": [ "-cp", - "/path/to/build/libs/mcp-server-0.0.1-all.jar", + "/path/to/smithy-java/examples/mcp-server/build/libs/mcp-server--all.jar", "software.amazon.smithy.java.example.server.mcp.ProxyMCPExample" ] } } } ``` - - - - diff --git a/examples/mcp-server/build.gradle.kts b/examples/mcp-server/build.gradle.kts index 2c93cba474..8a6d60dc2b 100644 --- a/examples/mcp-server/build.gradle.kts +++ b/examples/mcp-server/build.gradle.kts @@ -56,5 +56,9 @@ tasks.assemble { } java { - toolchain.languageVersion.set(JavaLanguageVersion.of(21)) + toolchain.languageVersion.set(JavaLanguageVersion.of(25)) +} + +tasks.withType() { + options.release.set(25) } diff --git a/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/MCPServerExample.java b/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/MCPServerExample.java index fe9a175970..343f368b7d 100644 --- a/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/MCPServerExample.java +++ b/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/MCPServerExample.java @@ -3,7 +3,7 @@ import software.amazon.smithy.java.example.server.mcp.operations.GetCodingStatistics; import software.amazon.smithy.java.example.server.mcp.operations.GetEmployeeDetails; import software.amazon.smithy.java.example.server.mcp.service.EmployeeService; -import software.amazon.smithy.java.mcp.server.McpServer; +import software.amazon.smithy.java.mcp.server.StdioMcpServer; public class MCPServerExample { @@ -13,7 +13,7 @@ public static void main(String[] args) { .addGetEmployeeDetailsOperation(new GetEmployeeDetails()) .build(); - var mcpServer = McpServer.builder() + var mcpServer = StdioMcpServer.builder() .stdio() .name("smithy-mcp-server") .addService("employee-mcp", service) @@ -22,8 +22,10 @@ public static void main(String[] args) { mcpServer.start(); try { - Thread.currentThread().join(); + mcpServer.awaitCompletion(); } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { mcpServer.shutdown(); } } diff --git a/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/ProxyMCPExample.java b/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/ProxyMCPExample.java index bdf75d9832..87f36dbd1a 100644 --- a/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/ProxyMCPExample.java +++ b/examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/ProxyMCPExample.java @@ -4,7 +4,7 @@ import software.amazon.smithy.java.example.server.mcp.operations.GetCodingStatistics; import software.amazon.smithy.java.example.server.mcp.operations.GetEmployeeDetails; import software.amazon.smithy.java.example.server.mcp.service.EmployeeService; -import software.amazon.smithy.java.mcp.server.McpServer; +import software.amazon.smithy.java.mcp.server.StdioMcpServer; import software.amazon.smithy.java.server.ProxyService; import software.amazon.smithy.java.server.Server; import software.amazon.smithy.model.Model; @@ -38,7 +38,7 @@ public static void main(String[] args) { .proxyEndpoint("http://localhost:8080") .build(); - var mcpServer = McpServer.builder() + var mcpServer = StdioMcpServer.builder() .stdio() .name("smithy-mcp-server") .addService("employee-mcp", mcpService) @@ -46,8 +46,10 @@ public static void main(String[] args) { mcpServer.start(); try { - Thread.currentThread().join(); + mcpServer.awaitCompletion(); } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { mcpServer.shutdown(); server.shutdown(); } diff --git a/mcp/mcp-schemas/model/main.smithy b/mcp/mcp-schemas/model/main.smithy index d0abb5dfc0..9d393c165d 100644 --- a/mcp/mcp-schemas/model/main.smithy +++ b/mcp/mcp-schemas/model/main.smithy @@ -55,6 +55,7 @@ structure InitializeResult with [BaseResult] { } structure Capabilities { + completions: Document logging: Document prompts: Prompts tools: Tools @@ -152,6 +153,10 @@ structure JsonPrimitiveSchema { /// JSON Schema format annotation (e.g., "date-time" for timestamps) format: String + + /// MCP HTTP parameter header suffix from smithy.ai#mcpHeader. + @jsonName("x-mcp-header") + mcpHeader: String } structure JsonDocumentSchema { diff --git a/mcp/mcp-server/README.md b/mcp/mcp-server/README.md index f0efe00961..b4becc8a91 100644 --- a/mcp/mcp-server/README.md +++ b/mcp/mcp-server/README.md @@ -5,3 +5,143 @@ > This module is not recommended for production use. Provides Model Context Protocol (MCP) server support for Smithy Java, enabling MCP server generation from Smithy models. + +## Creating a standard input/output server + +Generated Smithy services can be exposed directly: + +```java +var mcpServer = StdioMcpServer.builder() + .stdio() + .name("employee-server") + .version("1.0.0") + .addService("employees", employeeService) + .build(); + +mcpServer.start(); +mcpServer.awaitCompletion(); +``` + +For applications that need to share execution across transports, construct the +transport-independent engine separately: + +```java +var engine = McpEngine.builder() + .name("employee-server") + .addService("employees", employeeService) + .build(); + +var stdioServer = StdioMcpServer.builder() + .stdio() + .engine(engine) + .build(); +``` + +Builder-managed services and a prebuilt engine are mutually exclusive. + +## Architecture and extension points + +The implementation is split into a blocking, transport-independent `McpEngine`, +typed sealed `McpCall` and `McpOutcome` hierarchies, declarative per-version +protocol profiles, an immutable-snapshot source aggregator, and transport adapters: + +- `StdioMcpServer` exposes an engine over newline-delimited JSON-RPC and executes + requests on virtual threads. +- `McpHttpHandler` adapts decoded Streamable HTTP requests. +- `HttpMcpClient` and `StdioMcpClient` are blocking remote clients intended to + run naturally on virtual threads. +- `McpExtensionMethod` adds typed custom methods without modifying the built-in + protocol dispatch. +- `ExtensionMcpProtocol` is the open branch of the sealed `McpProtocol` + hierarchy for externally implemented protocol versions. + +Unsupported operations default to JSON-RPC method-not-found responses. A new +built-in protocol revision is added as one immutable method/feature declaration, +and the exhaustive version switch makes an incomplete registration fail at compile +time. + +## Adding a protocol + +Implement `ExtensionMcpProtocol` and override only the behavior that differs from +the defaults: + +```java +public final class FutureProtocol implements ExtensionMcpProtocol { + private static final McpProtocolId ID = McpProtocolId.of("2099-01-01"); + + @Override + public McpProtocolId id() { + return ID; + } + + @Override + public Set supportedMethods() { + return Set.of( + McpMethod.Standard.INITIALIZE, + McpMethod.Standard.PING, + McpMethod.Standard.TOOLS_LIST, + McpMethod.Standard.TOOLS_CALL); + } + + @Override + public McpProtocolFeatures features() { + return new McpProtocolFeatures(true, true, false, false, false); + } +} +``` + +Register it directly: + +```java +var engine = McpEngine.builder() + .addProtocol(new FutureProtocol()) + .build(); +``` + +Or publish it through Java's service-provider mechanism: + +```java +public final class FutureProtocolProvider implements McpProtocolProvider { + @Override + public Collection protocols() { + return List.of(new FutureProtocol()); + } +} +``` + +Register the provider class in: + +```text +META-INF/services/software.amazon.smithy.java.mcp.server.McpProtocolProvider +``` + +Built-in protocols, discovered providers, and builder registrations share one +immutable registry. Duplicate identifiers fail engine construction. This ensures +that upgrading to a release that implements a previously external protocol does +not silently change behavior. Use `overrideProtocol` only when replacement is +intentional: + +```java +var engine = McpEngine.builder() + .overrideProtocol(new FutureProtocol()) + .build(); +``` + +`discoverProtocols(false)` disables service-provider discovery. Programmatically +registered protocols remain enabled. + +Use `McpInterceptor` to observe or replace immutable calls and outcomes. Custom +method implementations use `McpExtensionMethod

` and are registered with +`McpEngine.Builder.addExtension`. Its outbound `encode` operation defaults to +`UnsupportedOperationException`, so inbound-only extensions implement only +decoding and execution. + +The module supports protocol revisions through `2026-07-28`. Run the official +Model Context Protocol conformance scenarios with: + +```console +./gradlew :mcp:mcp-server:conformance +``` + +The conformance task requires Node.js and invokes the pinned +`@modelcontextprotocol/conformance` package. diff --git a/mcp/mcp-server/build.gradle.kts b/mcp/mcp-server/build.gradle.kts index 4d760994be..868a6008ac 100644 --- a/mcp/mcp-server/build.gradle.kts +++ b/mcp/mcp-server/build.gradle.kts @@ -1,6 +1,7 @@ plugins { id("smithy-java.module-conventions") id("smithy-java.codegen-plugin-conventions") + `java-test-fixtures` } description = @@ -32,3 +33,20 @@ spotbugs { } addGenerateSrcsTask("software.amazon.smithy.java.mcp.server.utils.TestJavaCodegenRunner", null, null, "server") + +tasks.named("integ") { + useJUnitPlatform { + excludeTags("conformance") + } +} + +tasks.register("conformance") { + description = "Runs the official Model Context Protocol conformance scenarios" + group = "verification" + useJUnitPlatform { + includeTags("conformance") + } + testClassesDirs = sourceSets["it"].output.classesDirs + classpath = sourceSets["it"].runtimeClasspath + shouldRunAfter("integ") +} diff --git a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpConformanceTest.java b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpConformanceTest.java new file mode 100644 index 0000000000..0ee484ab11 --- /dev/null +++ b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpConformanceTest.java @@ -0,0 +1,302 @@ +/* + * 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.assertTrue; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.stream.Stream; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.json.JsonCodec; +import software.amazon.smithy.java.json.JsonSettings; +import software.amazon.smithy.java.mcp.conformance.model.TestCustomHeaderOutput; +import software.amazon.smithy.java.mcp.conformance.model.TestLoggingToolOutput; +import software.amazon.smithy.java.mcp.conformance.model.TestMissingCapabilityOutput; +import software.amazon.smithy.java.mcp.conformance.model.TestSimpleTextOutput; +import software.amazon.smithy.java.mcp.conformance.model.TestStreamingElicitationOutput; +import software.amazon.smithy.java.mcp.conformance.service.ConformanceService; +import software.amazon.smithy.java.mcp.conformance.service.TestCustomHeaderOperation; +import software.amazon.smithy.java.mcp.conformance.service.TestErrorHandlingOperation; +import software.amazon.smithy.java.mcp.conformance.service.TestLoggingToolOperation; +import software.amazon.smithy.java.mcp.conformance.service.TestMissingCapabilityOperation; +import software.amazon.smithy.java.mcp.conformance.service.TestSimpleTextOperation; +import software.amazon.smithy.java.mcp.conformance.service.TestStreamingElicitationOperation; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; + +@Tag("conformance") +class McpConformanceTest { + private static final String CONFORMANCE_VERSION = "0.2.0-alpha.11"; + private static final Duration PROCESS_TIMEOUT = Duration.ofMinutes(15); + private static final JsonCodec CODEC = JsonCodec.builder() + .settings(JsonSettings.builder() + .serializeTypeInDocuments(false) + .useJsonName(true) + .build()) + .build(); + + private static HttpServer directServer; + private static HttpServer proxyServer; + private static String directServerUrl; + private static String proxyServerUrl; + + @BeforeAll + static void startServer() throws IOException { + var service = createConformanceService(); + var directEngine = McpEngine.builder() + .services(Map.of("conformance", service)) + .name("smithy-java-conformance") + .version("1.0.0") + .interceptor(new RequiredCapabilityInterceptor()) + .build(); + directServer = startHttpServer(directEngine); + directServerUrl = serverUrl(directServer); + + var proxy = HttpMcpClient.builder() + .endpoint(directServerUrl) + .name("conformance-upstream") + .build(); + var proxyEngine = McpEngine.builder() + .remoteClients(List.of(proxy)) + .name("smithy-java-proxy-conformance") + .version("1.0.0") + .build(); + proxyServer = startHttpServer(proxyEngine); + proxyServerUrl = serverUrl(proxyServer); + } + + private static ConformanceService createConformanceService() { + return ConformanceService.builder() + .addTestCustomHeaderOperation( + (TestCustomHeaderOperation) (input, context) -> TestCustomHeaderOutput.builder() + .text("Custom header accepted: " + input.getValue()) + .build()) + .addTestErrorHandlingOperation( + (TestErrorHandlingOperation) (input, context) -> { + throw new RuntimeException("This tool intentionally returns an error for testing"); + }) + .addTestLoggingToolOperation( + (TestLoggingToolOperation) (input, context) -> TestLoggingToolOutput.builder() + .text("Logging completed.") + .build()) + .addTestMissingCapabilityOperation( + (TestMissingCapabilityOperation) (input, context) -> TestMissingCapabilityOutput.builder() + .text("Capability available.") + .build()) + .addTestSimpleTextOperation((TestSimpleTextOperation) (input, context) -> TestSimpleTextOutput.builder() + .text("This is a simple text response for testing.") + .build()) + .addTestStreamingElicitationOperation( + (TestStreamingElicitationOperation) (input, context) -> TestStreamingElicitationOutput.builder() + .text("Streaming elicitation completed.") + .build()) + .build(); + } + + private static HttpServer startHttpServer(McpEngine engine) throws IOException { + var requestHandler = McpHttpHandler.forLoopback(engine); + var server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext("/mcp", exchange -> handleHttpRequest(exchange, requestHandler)); + server.start(); + return server; + } + + private static String serverUrl(HttpServer server) { + return "http://127.0.0.1:" + server.getAddress().getPort() + "/mcp"; + } + + @AfterAll + static void stopServer() { + if (proxyServer != null) { + proxyServer.stop(0); + } + if (directServer != null) { + directServer.stop(0); + } + } + + @ParameterizedTest(name = "{0} requirements {1}") + @MethodSource("conformanceTopologies") + void passesOfficialConformanceRequirements( + String topology, + String protocolVersion, + String serverUrl + ) throws Exception { + var baseline = Path.of(McpConformanceTest.class + .getResource("/conformance-baseline-" + protocolVersion + ".yaml") + .toURI()) + .toString(); + var outputDirectory = Path.of( + "build", + "conformance-results", + topology + "-" + protocolVersion).toAbsolutePath(); + Files.createDirectories(outputDirectory); + var process = new ProcessBuilder( + "npx", + "--yes", + "@modelcontextprotocol/conformance@" + CONFORMANCE_VERSION, + "server", + "--url", + serverUrl, + "--requirements", + protocolVersion, + "--expected-failures", + baseline, + "--output-dir", + outputDirectory.toString(), + "--verbose") + .redirectErrorStream(true) + .start(); + + var output = CompletableFuture.supplyAsync(() -> readOutput(process)); + var exited = process.waitFor(PROCESS_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + if (!exited) { + process.destroyForcibly(); + } + + var commandOutput = output.get(10, TimeUnit.SECONDS); + assertTrue(exited, () -> "Conformance process timed out:\n" + commandOutput); + assertEquals(0, process.exitValue(), () -> "Conformance scenario failed:\n" + commandOutput); + } + + private static Stream conformanceTopologies() { + return Stream.of( + Arguments.of("direct", "2025-11-25", directServerUrl), + Arguments.of("proxy", "2025-11-25", proxyServerUrl), + Arguments.of("direct", "2026-07-28", directServerUrl), + Arguments.of("proxy", "2026-07-28", proxyServerUrl)); + } + + private static void handleHttpRequest(HttpExchange exchange, McpHttpHandler requestHandler) + throws IOException { + try (exchange) { + if (!"POST".equals(exchange.getRequestMethod())) { + exchange.sendResponseHeaders(405, -1); + return; + } + + var request = normalizeConformanceToolName( + CODEC.deserializeShape(exchange.getRequestBody().readAllBytes(), JsonRpcRequest.builder())); + var headers = normalizeConformanceToolNameHeader(request, exchange.getRequestHeaders()); + var response = requestHandler.handle(request, headers); + if (response.body() == null) { + exchange.sendResponseHeaders(response.statusCode(), -1); + return; + } + + var responseBytes = ByteBufferUtils.getBytes(CODEC.serialize(response.body())); + exchange.getResponseHeaders().set("Content-Type", "application/json"); + exchange.sendResponseHeaders(response.statusCode(), responseBytes.length); + exchange.getResponseBody().write(responseBytes); + } + } + + private static JsonRpcRequest normalizeConformanceToolName(JsonRpcRequest request) { + if (!"tools/call".equals(request.getMethod()) || request.getParams() == null) { + return request; + } + + var params = new HashMap<>(request.getParams().asStringMap()); + var name = params.get("name"); + if (name == null || !name.asString().contains("_")) { + return request; + } + + params.put("name", Document.of(toUpperCamelCase(name.asString()))); + return request.toBuilder().params(Document.of(params)).build(); + } + + private static Map> normalizeConformanceToolNameHeader( + JsonRpcRequest request, + Map> headers + ) { + var result = new HashMap>(); + headers.forEach((name, values) -> result.put(name, List.copyOf(values))); + if (!"tools/call".equals(request.getMethod())) { + return result; + } + + var nameHeader = headers.entrySet() + .stream() + .filter(entry -> entry.getKey().equalsIgnoreCase("mcp-name")) + .map(Map.Entry::getValue) + .filter(values -> !values.isEmpty()) + .map(List::getFirst) + .findFirst() + .orElse(null); + if (nameHeader == null || !nameHeader.contains("_")) { + return result; + } + + result.keySet().removeIf(name -> name.equalsIgnoreCase("mcp-name")); + result.put("mcp-name", List.of(toUpperCamelCase(nameHeader))); + return result; + } + + private static String toUpperCamelCase(String value) { + var result = new StringBuilder(value.length()); + var capitalizeNext = true; + for (var character : value.toCharArray()) { + if (character == '_') { + capitalizeNext = true; + } else if (capitalizeNext) { + result.append(Character.toUpperCase(character)); + capitalizeNext = false; + } else { + result.append(character); + } + } + return result.toString(); + } + + private static String readOutput(Process process) { + try (var output = new ByteArrayOutputStream()) { + process.getInputStream().transferTo(output); + return output.toString(StandardCharsets.UTF_8); + } catch (IOException e) { + return "Failed to read conformance process output: " + e.getMessage(); + } + } + + private static final class RequiredCapabilityInterceptor implements McpInterceptor { + @Override + public void readBeforeToolCall(McpToolExecutionContext hook) { + if (!"TestMissingCapability".equals(hook.call().name())) { + return; + } + + var capabilities = hook.call().metadata().clientCapabilities(); + if (capabilities == null || capabilities.getMember("sampling") == null) { + throw new McpProtocolException( + -32021, + "The sampling client capability is required", + Document.of(Map.of( + "requiredCapabilities", + Document.of(Map.of("sampling", Document.of(Map.of())))))); + } + } + } +} 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/StdioMcpServerIntegrationTest.java similarity index 96% rename from mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpServerIntegrationTest.java rename to mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/StdioMcpServerIntegrationTest.java index a4bacde7f4..126505fe74 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/StdioMcpServerIntegrationTest.java @@ -59,7 +59,7 @@ import tools.jackson.databind.ObjectMapper; import tools.jackson.databind.node.MissingNode; -class McpServerIntegrationTest { +class StdioMcpServerIntegrationTest { private static final JsonCodec CODEC = JsonCodec.builder() .settings(JsonSettings.builder() @@ -81,7 +81,7 @@ class McpServerIntegrationTest { private CalculateAreaOperation calculateAreaOperation; private int requestId = 0; private final Map outputSchemaCache = new HashMap<>(); - private ProtocolVersion currentProtocolVersion = null; + private KnownProtocolVersion currentProtocolVersion = null; @BeforeEach void init() { @@ -89,7 +89,7 @@ void init() { output = new TestOutputStream(); echoOperation = new McpEchoOperationImpl(); calculateAreaOperation = new CalculateAreaImpl(); - mcpServer = McpServer.builder() + mcpServer = StdioMcpServer.builder() .name("test-mcp") .addService("test-service", TestService.builder() @@ -114,7 +114,7 @@ void teardown() { // ========== Helper Methods ========== private void initializeLatestProtocol() { - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); } private Document getEchoFromResponse(JsonRpcResponse response) { @@ -133,7 +133,8 @@ 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 toolsNode = OBJECT_MAPPER.readTree(responseJson).path("result").path("tools"); + var toolNode = findToolNode(toolsNode, "McpEcho"); var inputSchemaNode = toolNode.path("inputSchema"); var outputSchemaNode = toolNode.path("outputSchema"); return new ToolSchemas( @@ -146,17 +147,22 @@ private ToolSchemas getCalculateAreaToolSchemas() { 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("CalculateArea")) { - var inputSchemaNode = toolNode.path("inputSchema"); - var outputSchemaNode = toolNode.path("outputSchema"); - return new ToolSchemas( - SCHEMA_FACTORY.getSchema(inputSchemaNode), - SCHEMA_FACTORY.getSchema(outputSchemaNode), - toolNode); + var toolNode = findToolNode(toolsNode, "CalculateArea"); + var inputSchemaNode = toolNode.path("inputSchema"); + var outputSchemaNode = toolNode.path("outputSchema"); + return new ToolSchemas( + SCHEMA_FACTORY.getSchema(inputSchemaNode), + SCHEMA_FACTORY.getSchema(outputSchemaNode), + toolNode); + } + + private JsonNode findToolNode(JsonNode tools, String name) { + for (var tool : tools) { + if (name.equals(tool.path("name").asString())) { + return tool; } } - throw new AssertionError("CalculateArea tool not found"); + throw new AssertionError(name + " tool not found"); } private String readRawResponse() { @@ -170,22 +176,27 @@ void testInitializeWithDefaultVersion() { write("initialize", Document.of(Map.of())); var response = read(); assertNotNull(response.getResult()); - assertEquals("2024-11-05", response.getResult().getMember("protocolVersion").asString()); + assertEquals("2025-03-26", response.getResult().getMember("protocolVersion").asString()); } @Test void testInitializeWithVersion2025_03_26() { - initializeWithProtocolVersion(ProtocolVersion.v2025_03_26.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_03_26); } @Test void testInitializeWithVersion2025_06_18() { - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); + } + + @Test + void testInitializeWithVersion2025_11_25() { + initializeWithProtocolVersion(KnownProtocolVersion.V2025_11_25); } @Test void testOutputSchemaNotPresentWithOlderProtocolVersion() { - initializeWithProtocolVersion(ProtocolVersion.v2025_03_26.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_03_26); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().getMember("tools").asList(); @@ -196,7 +207,7 @@ void testOutputSchemaNotPresentWithOlderProtocolVersion() { @Test void testOutputSchemaPresentWithVersion2025_06_18() { - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().getMember("tools").asList(); @@ -410,12 +421,12 @@ void testFloatingPointRoundTrip() { initializeLatestProtocol(); var echoInput = createEchoInput(Map.of( "floatValue", - Document.of(3.14f), + Document.of(1.25f), "doubleValue", - Document.of(2.718281828))); + Document.of(1.23456789))); var echo = getEchoFromResponse(callTool("McpEcho", echoInput)); - assertEquals(3.14f, echo.getMember("floatValue").asNumber().floatValue(), 0.001); - assertEquals(2.718281828, echo.getMember("doubleValue").asNumber().doubleValue(), 0.0000001); + assertEquals(1.25f, echo.getMember("floatValue").asNumber().floatValue(), 0.001); + assertEquals(1.23456789, echo.getMember("doubleValue").asNumber().doubleValue(), 0.0000001); } // ========== Big Number Tests ========== @@ -814,7 +825,7 @@ static Stream documentSchemaValidationTestCases() { // Number document (integer) Arguments.of("integer document", Document.of(42)), // Number document (double) - Arguments.of("double document", Document.of(3.14159)), + Arguments.of("double document", Document.of(1.23456)), // Boolean document (true) Arguments.of("boolean true document", Document.of(true)), // Boolean document (false) @@ -865,10 +876,9 @@ void testDocumentValidatesAgainstSchema(String description, Document documentVal 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) + var outputSchemaNode = findToolNode( + toolsResponseNode.path("result").path("tools"), + "McpEcho") .path("outputSchema"); // Create input with the document value @@ -1282,8 +1292,8 @@ void testInputFieldsAreCorrectlyDeserialized() { echoData.put("shortValue", Document.of(1000)); echoData.put("integerValue", Document.of(100000)); echoData.put("longValue", Document.of(9999999999L)); - echoData.put("floatValue", Document.of(3.14f)); - echoData.put("doubleValue", Document.of(2.718281828)); + echoData.put("floatValue", Document.of(1.25f)); + echoData.put("doubleValue", Document.of(1.23456789)); echoData.put("bigDecimalValue", Document.of("123.456")); echoData.put("bigIntegerValue", Document.of("123456789012345678901234567890")); echoData.put("blobValue", Document.of(base64Blob)); @@ -1308,8 +1318,8 @@ void testInputFieldsAreCorrectlyDeserialized() { assertEquals((short) 1000, echo.getShortValue().shortValue()); assertEquals(100000, echo.getIntegerValue().intValue()); assertEquals(9999999999L, echo.getLongValue().longValue()); - assertEquals(3.14f, echo.getFloatValue(), 0.001f); - assertEquals(2.718281828, echo.getDoubleValue(), 0.0000001); + assertEquals(1.25f, echo.getFloatValue(), 0.001f); + assertEquals(1.23456789, echo.getDoubleValue(), 0.0000001); // Verify big numbers assertEquals(new BigDecimal("123.456"), echo.getBigDecimalValue()); @@ -1350,10 +1360,9 @@ void testStructuredContentValidatesAgainstOutputSchema() { 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) + var outputSchemaNode = findToolNode( + toolsResponseNode.path("result").path("tools"), + "McpEcho") .path("outputSchema"); // Create comprehensive input @@ -1399,10 +1408,9 @@ void testAllTypesValidateAgainstOutputSchema() { 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) + var outputSchemaNode = findToolNode( + toolsResponseNode.path("result").path("tools"), + "McpEcho") .path("outputSchema"); // Create input with all types @@ -1451,7 +1459,7 @@ void testAllTypesValidateAgainstOutputSchema() { var callResponseNode = OBJECT_MAPPER.readTree(callResponseJson); var structuredContentNode = callResponseNode.path("result").path("structuredContent"); - assertFalse(structuredContentNode.isMissingNode(), "Missing structured content"); + assertFalse(structuredContentNode.isMissingNode(), "Missing structured content: " + callResponseJson); // Validate using Jackson-parsed JSON directly Schema schema = SCHEMA_FACTORY.getSchema(outputSchemaNode); @@ -1477,7 +1485,7 @@ void testUnknownTool() { @Test void testNoStructuredContentWithOlderProtocol() { - initializeWithProtocolVersion(ProtocolVersion.v2025_03_26.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_03_26); var response = callTool("McpEcho", createEchoInput(Map.of("stringValue", Document.of("test")))); // With older protocol, structuredContent should not be present assertNull(response.getResult().getMember("structuredContent")); @@ -2038,6 +2046,48 @@ void testCalculateAreaWithCircle() { assertEquals(expectedArea, result.getMember("area").asNumber().doubleValue(), 0.001); } + @Test + void testToolExecutionFailureReturnsMcpErrorResult() { + mcpServer.shutdown().join(); + input = new TestInputStream(); + output = new TestOutputStream(); + mcpServer = StdioMcpServer.builder() + .name("test-mcp") + .addService("test-service", + TestService.builder() + .addCalculateAreaOperation( + (CalculateAreaOperation) (input, context) -> { + throw new RuntimeException("tool execution failed"); + }) + .addMcpEchoOperation(echoOperation) + .build()) + .input(input) + .output(output) + .build(); + mcpServer.start(); + + initializeWithProtocolVersion(KnownProtocolVersion.V2025_11_25); + var circle = Document.of(Map.of("circle", Document.of(Map.of("radius", Document.of(5))))); + write("tools/call", + Document.of(Map.of( + "name", + Document.of("CalculateArea"), + "arguments", + Document.of(Map.of("oneOfInput", circle))))); + + var response = read(); + assertNull(response.getError()); + assertTrue(response.getResult().getMember("isError").asBoolean()); + assertEquals( + "tool execution failed", + response.getResult() + .getMember("content") + .asList() + .getFirst() + .getMember("text") + .asString()); + } + @Test void testCalculateAreaWithSquare() { initializeLatestProtocol(); @@ -2331,7 +2381,7 @@ void testRecursiveOneOfSchemaTerminatesWithoutInfiniteLoop() { // ========== Helper Methods ========== - private void initializeWithProtocolVersion(ProtocolVersion protocolVersion) { + private void initializeWithProtocolVersion(KnownProtocolVersion protocolVersion) { this.currentProtocolVersion = protocolVersion; var params = Document.of(Map.of("protocolVersion", Document.of(protocolVersion.identifier()))); write("initialize", params); @@ -2364,7 +2414,7 @@ private void cacheToolSchemas() { private void validateStructuredContent(String toolName, String responseJson) { // Only validate for protocol versions that support structured content (v2025_06_18+) boolean supportsStructuredContent = currentProtocolVersion != null - && currentProtocolVersion.compareTo(ProtocolVersion.v2025_06_18.INSTANCE) >= 0; + && currentProtocolVersion.compareTo(KnownProtocolVersion.V2025_06_18) >= 0; if (!supportsStructuredContent) { return; // Skip validation for older protocols @@ -2380,7 +2430,10 @@ private void validateStructuredContent(String toolName, String responseJson) { // Assert structured content IS present for compatible protocols assertFalse(structuredContentNode.isMissingNode(), - "structuredContent should be present for protocol version " + currentProtocolVersion.identifier()); + "structuredContent should be present for protocol version " + + currentProtocolVersion.identifier() + + ": " + + responseJson); // Validate against schema var schema = outputSchemaCache.get(toolName); diff --git a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java b/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java deleted file mode 100644 index 11f049d295..0000000000 --- a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java +++ /dev/null @@ -1,88 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import java.io.ByteArrayOutputStream; -import java.io.OutputStream; -import java.nio.charset.StandardCharsets; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingQueue; - -final class TestOutputStream extends OutputStream { - private final BlockingQueue lines = new LinkedBlockingQueue<>(); - private final ByteArrayOutputStream baos = new ByteArrayOutputStream(); - - @Override - public void write(int b) { - baos.write(b); - if (b == '\n') { - lines.add(baos.toString(StandardCharsets.UTF_8)); - baos.reset(); - } - } - - @Override - public void write(byte[] b, int off, int len) { - int rem = len; - int pos = off; - while (rem > 0) { - int nl = find(b, pos, pos + rem, (byte) '\n'); - if (nl == -1) { - baos.write(b, pos, rem); - return; - } else { - // Include the newline character in what we write - int toWrite = nl - pos + 1; - baos.write(b, pos, toWrite); - lines.add(baos.toString(StandardCharsets.UTF_8)); - baos.reset(); - rem -= toWrite; - pos += toWrite; - } - } - } - - private static int find(byte[] arr, int start, int end, byte b) { - if (start >= end || end > arr.length) { - throw new IllegalArgumentException(); - } - for (int i = start; i < end; i++) { - if (arr[i] == b) { - return i; - } - } - return -1; - } - - String read() { - try { - return lines.take(); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - } - - boolean hasOutput() { - return !lines.isEmpty(); - } - - void assertNoOutput() { - assertNoOutput(50); - } - - void assertNoOutput(long waitMillis) { - // Wait briefly to allow any potential response to be written - try { - Thread.sleep(waitMillis); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - // Verify no output was produced - if (hasOutput()) { - throw new AssertionError("Expected no output but got : " + read()); - } - } -} diff --git a/mcp/mcp-server/src/it/resources/META-INF/smithy/conformance.smithy b/mcp/mcp-server/src/it/resources/META-INF/smithy/conformance.smithy new file mode 100644 index 0000000000..017bc174b1 --- /dev/null +++ b/mcp/mcp-server/src/it/resources/META-INF/smithy/conformance.smithy @@ -0,0 +1,82 @@ +$version: "2" + +namespace software.amazon.smithy.java.mcp.conformance + +use smithy.ai#mcpHeader +use smithy.ai#prompts + +@prompts({ + test_simple_prompt: { + description: "A simple prompt without arguments" + template: "This is a simple prompt for testing." + } + test_prompt_with_arguments: { + description: "A prompt with required arguments" + template: "Prompt with arguments: arg1='{{arg1}}', arg2='{{arg2}}'" + arguments: TestPromptArguments + } +}) +service ConformanceService { + operations: [ + TestCustomHeader, + TestErrorHandling, + TestMissingCapability, + TestStreamingElicitation, + TestLoggingTool, + TestSimpleText + ] +} + +operation TestCustomHeader { + input: TestCustomHeaderInput + output: ConformanceTextOutput +} + +operation TestErrorHandling { + input := {} + output: ConformanceTextOutput +} + +operation TestMissingCapability { + input := {} + output: ConformanceTextOutput +} + +operation TestStreamingElicitation { + input := {} + output: ConformanceTextOutput +} + +operation TestLoggingTool { + input := {} + output: ConformanceTextOutput +} + +operation TestSimpleText { + input := {} + output: ConformanceTextOutput +} + +structure ConformanceTextOutput { + @required + text: String + + // Keeps the prompt argument schema reachable from generated runtime schemas. + promptArguments: TestPromptArguments +} + +structure TestPromptArguments { + @required + @documentation("First test argument") + arg1: String + + @required + @documentation("Second test argument") + arg2: String +} + +structure TestCustomHeaderInput { + @required + @mcpHeader("test-value") + value: String +} diff --git a/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy b/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy index ac05d50e72..beea8558e6 100644 --- a/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy +++ b/mcp/mcp-server/src/it/resources/META-INF/smithy/main.smithy @@ -7,7 +7,7 @@ use smithy.mcp#oneOf service TestService { operations: [ McpEcho, - CalculateArea + CalculateArea ] } diff --git a/mcp/mcp-server/src/it/resources/META-INF/smithy/manifest b/mcp/mcp-server/src/it/resources/META-INF/smithy/manifest index d2ade8b506..7206916980 100644 --- a/mcp/mcp-server/src/it/resources/META-INF/smithy/manifest +++ b/mcp/mcp-server/src/it/resources/META-INF/smithy/manifest @@ -1 +1,2 @@ -main.smithy \ No newline at end of file +main.smithy +conformance.smithy diff --git a/mcp/mcp-server/src/it/resources/conformance-baseline-2025-11-25.yaml b/mcp/mcp-server/src/it/resources/conformance-baseline-2025-11-25.yaml new file mode 100644 index 0000000000..48fed5c074 --- /dev/null +++ b/mcp/mcp-server/src/it/resources/conformance-baseline-2025-11-25.yaml @@ -0,0 +1,22 @@ +server: + - tools-call-image + - tools-call-audio + - tools-call-embedded-resource + - tools-call-mixed-content + - tools-call-with-logging + - tools-call-with-progress + - tools-call-sampling + - tools-call-elicitation + - elicitation-sep1034-defaults + - elicitation-sep1330-enums + # The runner treats this zero-check scenario as failed during baseline evaluation. + - server-sse-multiple-streams + - resources-list + - resources-read-text + - resources-read-binary + - resources-templates-read + - resources-subscribe + - resources-unsubscribe + - prompts-get-embedded-resource + - prompts-get-with-image + - json-schema-2020-12 diff --git a/mcp/mcp-server/src/it/resources/conformance-baseline-2026-07-28.yaml b/mcp/mcp-server/src/it/resources/conformance-baseline-2026-07-28.yaml new file mode 100644 index 0000000000..4d0c088c7a --- /dev/null +++ b/mcp/mcp-server/src/it/resources/conformance-baseline-2026-07-28.yaml @@ -0,0 +1,37 @@ +server: + - tools-call-image + - tools-call-audio + - tools-call-embedded-resource + - tools-call-mixed-content + - tools-call-with-progress + - resources-list + - resources-read-text + - resources-read-binary + - resources-templates-read + # These contain skipped internal legs that baseline evaluation treats as failures. + - sep-2164-resource-not-found + - prompts-get-embedded-resource + - prompts-get-with-image + - caching + - input-required-result-basic-elicitation + - input-required-result-basic-sampling + - input-required-result-basic-list-roots + - input-required-result-request-state + - input-required-result-multiple-input-requests + - input-required-result-multi-round + - input-required-result-missing-input-response + - input-required-result-non-tool-request + - input-required-result-result-type + - input-required-result-tampered-state + - input-required-result-capability-check + - input-required-result-ignore-extra-params + - tasks-lifecycle + - tasks-capability-negotiation + - tasks-wire-fields + - tasks-request-state-removal + - tasks-mrtr-input + - tasks-request-headers + - tasks-dispatch-and-envelope + - tasks-required-task-error + - tasks-mrtr-composition + - json-schema-2020-12 diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocol.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocol.java new file mode 100644 index 0000000000..af77c32005 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocol.java @@ -0,0 +1,67 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import software.amazon.smithy.java.core.serde.document.Document; + +/** + * Declarative behavior for one built-in MCP protocol version. + */ +record BuiltInProtocol( + KnownProtocolVersion version, + Set supportedMethods, + McpProtocolFeatures features) implements McpProtocol { + + BuiltInProtocol { + supportedMethods = Set.copyOf(supportedMethods); + } + + @Override + public McpProtocolId id() { + return version.id(); + } + + @Override + public Document decorateResult( + Document result, + McpMethod method, + McpServerIdentity serverIdentity + ) { + if (!features.statelessResults()) { + return result; + } + + var members = new HashMap<>(result.asStringMap()); + members.put("resultType", Document.of("complete")); + + var meta = members.containsKey("_meta") + ? new HashMap<>(members.get("_meta").asStringMap()) + : new HashMap(); + meta.put(McpWireNames.SERVER_INFO, + Document.of(Map.of( + "name", + Document.of(serverIdentity.name()), + "version", + Document.of(serverIdentity.version())))); + members.put("_meta", Document.of(meta)); + + if (switch (method) { + case McpMethod.Standard.SERVER_DISCOVER, + McpMethod.Standard.TOOLS_LIST, + McpMethod.Standard.PROMPTS_LIST -> + true; + default -> false; + }) { + members.put("ttlMs", Document.of(0)); + members.put("cacheScope", Document.of("private")); + } + return Document.of(members); + } + +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocols.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocols.java new file mode 100644 index 0000000000..0239b7928e --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocols.java @@ -0,0 +1,83 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.List; +import java.util.Set; + +final class BuiltInProtocols { + private static final McpProtocolFeatures LEGACY = + new McpProtocolFeatures(false, false, false, false, false); + private static final McpProtocolFeatures ANNOTATIONS = + new McpProtocolFeatures(false, true, false, false, false); + private static final McpProtocolFeatures STRUCTURED_OUTPUT = + new McpProtocolFeatures(true, true, false, false, false); + private static final McpProtocolFeatures STATELESS = + new McpProtocolFeatures(true, true, true, true, true); + + private static final Set LEGACY_METHODS = Set.of( + McpMethod.Standard.INITIALIZE, + McpMethod.Standard.PING, + McpMethod.Standard.PROMPTS_LIST, + McpMethod.Standard.PROMPTS_GET, + McpMethod.Standard.COMPLETION_COMPLETE, + McpMethod.Standard.LOGGING_SET_LEVEL, + McpMethod.Standard.TOOLS_LIST, + McpMethod.Standard.TOOLS_CALL, + McpMethod.Standard.NOTIFICATIONS_INITIALIZED, + McpMethod.Standard.NOTIFICATIONS_TOOLS_LIST_CHANGED); + private static final Set STATELESS_METHODS = Set.of( + McpMethod.Standard.SERVER_DISCOVER, + McpMethod.Standard.PROMPTS_LIST, + McpMethod.Standard.PROMPTS_GET, + McpMethod.Standard.COMPLETION_COMPLETE, + McpMethod.Standard.TOOLS_LIST, + McpMethod.Standard.TOOLS_CALL); + + private static final BuiltInProtocol V2024_11_05 = new BuiltInProtocol( + KnownProtocolVersion.V2024_11_05, + LEGACY_METHODS, + LEGACY); + private static final BuiltInProtocol V2025_03_26 = new BuiltInProtocol( + KnownProtocolVersion.V2025_03_26, + LEGACY_METHODS, + ANNOTATIONS); + private static final BuiltInProtocol V2025_06_18 = new BuiltInProtocol( + KnownProtocolVersion.V2025_06_18, + LEGACY_METHODS, + STRUCTURED_OUTPUT); + private static final BuiltInProtocol V2025_11_25 = new BuiltInProtocol( + KnownProtocolVersion.V2025_11_25, + LEGACY_METHODS, + STRUCTURED_OUTPUT); + private static final BuiltInProtocol V2026_07_28 = new BuiltInProtocol( + KnownProtocolVersion.V2026_07_28, + STATELESS_METHODS, + STATELESS); + + private static final List ALL = List.of( + V2026_07_28, + V2025_11_25, + V2025_06_18, + V2025_03_26, + V2024_11_05); + + private BuiltInProtocols() {} + + static BuiltInProtocol protocol(KnownProtocolVersion version) { + return switch (version) { + case V2024_11_05 -> V2024_11_05; + case V2025_03_26 -> V2025_03_26; + case V2025_06_18 -> V2025_06_18; + case V2025_11_25 -> V2025_11_25; + case V2026_07_28 -> V2026_07_28; + }; + } + + static List all() { + return ALL; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ExtensionMcpProtocol.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ExtensionMcpProtocol.java new file mode 100644 index 0000000000..594a357d9a --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ExtensionMcpProtocol.java @@ -0,0 +1,14 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Open SPI branch for externally implemented MCP protocols. + */ +@SmithyUnstableApi +public non-sealed interface ExtensionMcpProtocol extends McpProtocol {} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpProxy.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpClient.java similarity index 61% rename from mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpProxy.java rename to mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpClient.java index 97f85e40b3..414cf56420 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpProxy.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpClient.java @@ -5,10 +5,15 @@ package software.amazon.smithy.java.mcp.server; +import java.io.BufferedReader; +import java.io.InputStreamReader; import java.net.URI; import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.util.concurrent.CompletableFuture; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; import software.amazon.smithy.java.auth.api.Signer; import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; @@ -21,25 +26,20 @@ import software.amazon.smithy.java.http.api.HeaderName; import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.http.api.ModifiableHttpRequest; import software.amazon.smithy.java.io.ByteBufferUtils; import software.amazon.smithy.java.io.datastream.DataStream; -import software.amazon.smithy.java.json.JsonCodec; -import software.amazon.smithy.java.json.JsonSettings; import software.amazon.smithy.java.logging.InternalLogger; 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.ToolInfo; import software.amazon.smithy.utils.SmithyUnstableApi; @SmithyUnstableApi -public final class HttpMcpProxy extends McpServerProxy { - private static final InternalLogger LOG = InternalLogger.getLogger(HttpMcpProxy.class); - private static final JsonCodec JSON_CODEC = JsonCodec.builder() - .settings(JsonSettings.builder().serializeTypeInDocuments(false).useJsonName(true).build()) - .build(); - - private static final HeaderName MCP_PROTOCOL_VERSION = HeaderName.of("mcp-protocol-version"); - private static final HeaderName MCP_SESSION_ID = HeaderName.of("mcp-session-id"); +public final class HttpMcpClient extends McpRemoteClient { + private static final InternalLogger LOG = InternalLogger.getLogger(HttpMcpClient.class); + private static final int UPSTREAM_HTTP_ERROR_CODE = -32000; private final ClientTransport transport; private final URI endpoint; @@ -49,9 +49,11 @@ public final class HttpMcpProxy extends McpServerProxy { private final IdentityResolver identityResolver; private final Context signerContext; private final Duration timeout; + private final AtomicReference>> toolHeaderParameters = + new AtomicReference<>(Map.of()); private volatile String sessionId; - private HttpMcpProxy(Builder builder) { + private HttpMcpClient(Builder builder) { this.transport = builder.transport != null ? builder.transport : new JavaHttpClientTransport(); this.endpoint = URI.create(builder.endpoint); this.name = builder.name != null ? builder.name : sanitizeName(endpoint.getHost()); @@ -119,7 +121,7 @@ public Builder timeout(Duration timeout) { return this; } - public HttpMcpProxy build() { + public HttpMcpClient build() { if (endpoint == null || endpoint.isEmpty()) { throw new IllegalArgumentException("Endpoint must be provided"); } @@ -135,7 +137,7 @@ public HttpMcpProxy build() { throw new IllegalArgumentException( "authScheme must be provided when identityResolver is set"); } - return new HttpMcpProxy(this); + return new HttpMcpClient(this); } } @@ -144,24 +146,47 @@ public static Builder builder() { } @Override - public CompletableFuture rpc(JsonRpcRequest request) { + public List listTools() { + var tools = super.listTools(); + var updatedMappings = new HashMap>(); + for (var tool : tools) { + var mappings = McpHttpBinding.headerParameters(tool); + if (!mappings.isEmpty()) { + updatedMappings.put(tool.getName(), mappings); + } + } + toolHeaderParameters.set(Map.copyOf(updatedMappings)); + return tools; + } + + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { try { - byte[] body = JSON_CODEC.serializeToString(request).getBytes(StandardCharsets.UTF_8); + byte[] body = ByteBufferUtils.getBytes(McpJson.CODEC.serialize(request)); LOG.trace("Sending HTTP request to {}", endpoint); - String protocolVersionHeader = getProtocolVersion().identifier(); + var protocol = requestProtocol(request); var requestBuilder = HttpRequest.create() .setUri(endpoint) .setMethod("POST") .addHeader(HeaderName.CONTENT_TYPE, "application/json") .addHeader(HeaderName.ACCEPT, "application/json, text/event-stream") - .addHeader(MCP_PROTOCOL_VERSION, protocolVersionHeader); + .addHeader(McpHttpBinding.PROTOCOL_VERSION, protocol.id().identifier()); + + if (McpHttpBinding.usesMethodHeaders(protocol)) { + requestBuilder.addHeader(McpHttpBinding.METHOD, request.getMethod()); + var requestName = McpHttpBinding.requestName(request); + if (requestName != null) { + requestBuilder.addHeader(McpHttpBinding.NAME, requestName); + } + addParameterHeaders(requestBuilder, request, requestName); + } // Include session ID if we have one String currentSessionId = sessionId; if (currentSessionId != null) { - requestBuilder.addHeader(MCP_SESSION_ID, currentSessionId); + requestBuilder.addHeader(McpHttpBinding.SESSION_ID, currentSessionId); LOG.debug("Including session ID in request: method={}, sessionId={}", request.getMethod(), currentSessionId); @@ -186,7 +211,7 @@ public CompletableFuture rpc(JsonRpcRequest request) { LOG.trace("Received HTTP response with status: {}", response.statusCode()); // Extract and store session ID from response only during initialize - if ("initialize".equals(request.getMethod())) { + if (McpHttpBinding.isInitialize(request)) { String responseSessionId = response.headers().firstValue("Mcp-Session-Id"); if (responseSessionId != null) { sessionId = responseSessionId; @@ -202,20 +227,69 @@ public CompletableFuture rpc(JsonRpcRequest request) { } if (response.statusCode() < 200 || response.statusCode() >= 300) { - return CompletableFuture.completedFuture(handleErrorResponse(response)); + return handleErrorResponse(response, request); } // Check if response is SSE String contentType = response.body().contentType(); - if ("text/event-stream".equals(contentType)) { - return CompletableFuture.completedFuture(parseSseResponse(response, request)); + if (contentType != null && contentType.startsWith("text/event-stream")) { + return parseSseResponse(response, request); } - return CompletableFuture.completedFuture(JsonRpcResponse.builder() - .deserialize(JSON_CODEC.createDeserializer(response.body().asByteBuffer())) - .build()); + var responseBytes = ByteBufferUtils.getBytes(response.body().asByteBuffer()); + if (responseBytes.length == 0) { + return null; + } + return JsonRpcResponse.builder() + .deserialize(McpJson.CODEC.createDeserializer(responseBytes)) + .build(); } catch (Exception e) { - return CompletableFuture.failedFuture(e); + throw new McpRemoteException("HTTP MCP exchange failed", e); + } + } + + private McpProtocol requestProtocol(JsonRpcRequest request) { + var params = request.getParams(); + var meta = params == null ? null : params.getMember("_meta"); + var requestedVersion = meta == null ? null : meta.getMember(McpWireNames.PROTOCOL_VERSION); + var selected = protocol(); + if (requestedVersion == null || requestedVersion.asString().equals(selected.id().identifier())) { + return selected; + } + var parsed = ProtocolVersion.parse(requestedVersion.asString()); + if (parsed instanceof KnownProtocolVersion known) { + return BuiltInProtocols.protocol(known); + } + throw new McpRemoteException("Unregistered MCP protocol: " + parsed.identifier()); + } + + private String stringMember(Document document, String name) { + return McpHttpBinding.stringMember(document, name); + } + + private void addParameterHeaders( + ModifiableHttpRequest requestBuilder, + JsonRpcRequest request, + String toolName + ) { + if (!McpHttpBinding.isToolCall(request) || toolName == null) { + return; + } + + var mappings = toolHeaderParameters.get().get(toolName); + var params = request.getParams(); + var arguments = params == null ? null : params.getMember("arguments"); + if (mappings == null || arguments == null) { + return; + } + + for (var entry : mappings.entrySet()) { + var value = stringMember(arguments, entry.getKey()); + if (value != null) { + requestBuilder.addHeader( + HeaderName.of("Mcp-Param-" + entry.getValue()), + McpHttpBinding.encodeParameter(value)); + } } } @@ -234,99 +308,25 @@ private HttpRequest signWithAuthScheme(HttpRequest request) } private JsonRpcResponse parseSseResponse(HttpResponse response, JsonRpcRequest request) { - try { - byte[] bodyBytes = ByteBufferUtils.getBytes(response.body().asByteBuffer()); - String sseContent = new String(bodyBytes, StandardCharsets.UTF_8); - - JsonRpcResponse finalResponse = null; - Iterable lines = sseContent.lines()::iterator; - StringBuilder dataBuffer = new StringBuilder(); - - for (String line : lines) { - if (line.startsWith("data:")) { - var value = line.substring(5); - dataBuffer.append(value.startsWith(" ") ? value.substring(1) : value); - } else if (line.trim().isEmpty() && !dataBuffer.isEmpty()) { - // End of an SSE event - String jsonData = dataBuffer.toString().trim(); - dataBuffer.setLength(0); - - if (jsonData.isEmpty()) { - continue; - } - - try { - // Parse JSON once into Document - Document jsonDocument = JSON_CODEC.createDeserializer(jsonData.getBytes(StandardCharsets.UTF_8)) - .readDocument(); - - // Check if it's a notification by checking for top-level "id" field - // Notifications have "method" but no "id", responses have "id" - if (isNotification(jsonDocument)) { - // This is a notification - convert Document to JsonRpcRequest and forward - JsonRpcRequest notification = jsonDocument.asShape(JsonRpcRequest.builder()); - LOG.debug("Received notification from SSE stream: method={}", notification.getMethod()); - notify(notification); - } else { - // This is a response - convert Document to JsonRpcResponse - finalResponse = jsonDocument.asShape(JsonRpcResponse.builder()); - } - } catch (Exception e) { - LOG.warn("Failed to parse SSE message: {}", jsonData, e); - } - } - } - - // Process any remaining data in buffer (in case stream doesn't end with empty line) - if (!dataBuffer.isEmpty()) { - String jsonData = dataBuffer.toString().trim(); - if (!jsonData.isEmpty()) { - try { - // Parse JSON once into Document - Document jsonDocument = JSON_CODEC.createDeserializer(jsonData.getBytes(StandardCharsets.UTF_8)) - .readDocument(); - - // Check if it's a notification by checking for top-level "id" field - // Notifications have "method" but no "id", responses have "id" - if (isNotification(jsonDocument)) { - JsonRpcRequest notification = JsonRpcRequest.builder() - .deserialize(jsonDocument.createDeserializer()) - .build(); - LOG.debug("Received notification from remaining SSE buffer: method={}", - notification.getMethod()); - notify(notification); - } else { - JsonRpcResponse message = JsonRpcResponse.builder() - .deserialize(jsonDocument.createDeserializer()) - .build(); - - if (message.getId() == null) { - notify(JsonRpcRequest.builder() - .jsonrpc("2.0") - .method("notifications/unknown") - .build()); - } else { - finalResponse = message; - } - } - } catch (Exception e) { - LOG.warn("Failed to parse remaining SSE message: {}", jsonData, e); + try (var input = response.body().asInputStream(); + var reader = new BufferedReader(new InputStreamReader(input, StandardCharsets.UTF_8))) { + var data = new StringBuilder(); + String line; + while ((line = reader.readLine()) != null) { + if (line.isEmpty()) { + var result = processSseEvent(data); + if (result != null) { + return result; } + } else if (line.startsWith("data:")) { + var value = line.substring(5); + data.append(value.startsWith(" ") ? value.substring(1) : value).append('\n'); } } - - if (finalResponse == null) { - return JsonRpcResponse.builder() - .jsonrpc("2.0") - .id(request.getId()) - .error(JsonRpcErrorResponse.builder() - .code(-32001) - .message("SSE parsing error: No final response found in stream") - .build()) - .build(); + var result = processSseEvent(data); + if (result != null) { + return result; } - - return finalResponse; } catch (Exception e) { LOG.error("Error parsing SSE response", e); return JsonRpcResponse.builder() @@ -338,27 +338,62 @@ private JsonRpcResponse parseSseResponse(HttpResponse response, JsonRpcRequest r .build()) .build(); } + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .error(JsonRpcErrorResponse.builder() + .code(-32001) + .message("SSE parsing error: No final response found in stream") + .build()) + .build(); + } + + private JsonRpcResponse processSseEvent(StringBuilder data) { + if (data.isEmpty()) { + return null; + } + var jsonData = data.toString().stripTrailing(); + data.setLength(0); + if (jsonData.isEmpty()) { + return null; + } + + try { + var document = McpJson.CODEC.createDeserializer(jsonData.getBytes(StandardCharsets.UTF_8)) + .readDocument(); + if (isNotification(document)) { + var notification = document.asShape(JsonRpcRequest.builder()); + LOG.debug("Received notification from SSE stream: method={}", notification.getMethod()); + notify(notification); + return null; + } + return document.asShape(JsonRpcResponse.builder()); + } catch (RuntimeException e) { + LOG.warn("Failed to parse SSE message: {}", jsonData, e); + return null; + } } - private JsonRpcResponse handleErrorResponse(HttpResponse response) { + private JsonRpcResponse handleErrorResponse(HttpResponse response, JsonRpcRequest request) { long contentLength = response.body().contentLength(); String errorMessage = "HTTP " + response.statusCode(); - if (contentLength > 0) { + if (contentLength != 0) { String contentType = response.body().contentType(); byte[] bodyBytes = ByteBufferUtils.getBytes(response.body().asByteBuffer()); - if ("application/json".equals(contentType)) { + if (contentType != null && contentType.startsWith("application/json")) { try { return JsonRpcResponse.builder() - .deserialize(JSON_CODEC.createDeserializer(bodyBytes)) + .deserialize(McpJson.CODEC.createDeserializer(bodyBytes)) .build(); } catch (Exception e) { LOG.warn("Failed to deserialize JSON error response", e); return JsonRpcResponse.builder() .jsonrpc("2.0") + .id(request.getId()) .error(JsonRpcErrorResponse.builder() - .code(response.statusCode()) + .code(UPSTREAM_HTTP_ERROR_CODE) .message("HTTP " + response.statusCode() + ": Invalid JSON response") .build()) .build(); @@ -372,8 +407,9 @@ private JsonRpcResponse handleErrorResponse(HttpResponse response) { return JsonRpcResponse.builder() .jsonrpc("2.0") + .id(request.getId()) .error(JsonRpcErrorResponse.builder() - .code(response.statusCode()) + .code(UPSTREAM_HTTP_ERROR_CODE) .message(errorMessage) .build()) .build(); @@ -386,10 +422,9 @@ public void start() { } @Override - public CompletableFuture shutdown() { + public void close() { // HTTP client doesn't need explicit shutdown LOG.debug("HTTP MCP proxy shutdown for endpoint: {}", endpoint); - return CompletableFuture.completedFuture(null); } @Override diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/KnownProtocolVersion.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/KnownProtocolVersion.java new file mode 100644 index 0000000000..da835055a6 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/KnownProtocolVersion.java @@ -0,0 +1,73 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Protocol versions implemented by the server. + * + *

The declaration order is chronological. Code that varies by protocol behavior + * must use {@link BuiltInProtocols#protocol(KnownProtocolVersion)} rather than comparing + * versions. + */ +@SmithyUnstableApi +public enum KnownProtocolVersion implements ProtocolVersion { + V2024_11_05("2024-11-05"), + V2025_03_26("2025-03-26"), + V2025_06_18("2025-06-18"), + V2025_11_25("2025-11-25"), + V2026_07_28("2026-07-28"); + + private static final Map BY_IDENTIFIER; + private static final List SUPPORTED_IDENTIFIERS; + + static { + var byIdentifier = new LinkedHashMap(); + for (var version : values()) { + if (byIdentifier.put(version.identifier, version) != null) { + throw new IllegalStateException("Duplicate MCP protocol version: " + version.identifier); + } + } + BY_IDENTIFIER = Collections.unmodifiableMap(byIdentifier); + var versions = values(); + var supported = new ArrayList(versions.length); + for (int index = versions.length - 1; index >= 0; index--) { + supported.add(versions[index].identifier()); + } + SUPPORTED_IDENTIFIERS = List.copyOf(supported); + } + + private final String identifier; + private final McpProtocolId id; + + KnownProtocolVersion(String identifier) { + this.identifier = identifier; + this.id = McpProtocolId.of(identifier); + } + + @Override + public String identifier() { + return identifier; + } + + public McpProtocolId id() { + return id; + } + + static KnownProtocolVersion fromIdentifier(String identifier) { + return BY_IDENTIFIER.get(identifier); + } + + static List supportedIdentifiers() { + return SUPPORTED_IDENTIFIERS; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCall.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCall.java new file mode 100644 index 0000000000..e111518b30 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCall.java @@ -0,0 +1,217 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Map; +import java.util.Objects; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * A decoded, typed MCP call. + * + *

Dynamic Smithy tool arguments and extension payloads remain documents because + * their schemas are selected at runtime. Standard protocol parameters are represented + * explicitly by records. + */ +@SmithyUnstableApi +public sealed interface McpCall permits + McpCall.Initialize, + McpCall.Ping, + McpCall.Discover, + McpCall.ListTools, + McpCall.CallTool, + McpCall.ListPrompts, + McpCall.GetPrompt, + McpCall.Complete, + McpCall.SetLogLevel, + McpCall.ReadResource, + McpCall.Notification, + McpCall.ExtensionCall, + McpCall.UnknownCall { + + Document id(); + + McpMethod method(); + + McpMetadata metadata(); + + record Initialize( + Document id, + ProtocolVersion requestedVersion, + Document clientInfo, + Document capabilities, + McpMetadata metadata) implements McpCall { + public Initialize { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.INITIALIZE; + } + } + + record Ping(Document id, McpMetadata metadata) implements McpCall { + public Ping { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.PING; + } + } + + record Discover(Document id, McpMetadata metadata) implements McpCall { + public Discover { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.SERVER_DISCOVER; + } + } + + record ListTools(Document id, String cursor, McpMetadata metadata) implements McpCall { + public ListTools { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.TOOLS_LIST; + } + } + + record CallTool(Document id, String name, Document arguments, McpMetadata metadata) implements McpCall { + public CallTool { + Objects.requireNonNull(name, "name"); + arguments = arguments == null ? Document.of(Map.of()) : arguments; + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.TOOLS_CALL; + } + } + + record ListPrompts(Document id, String cursor, McpMetadata metadata) implements McpCall { + public ListPrompts { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.PROMPTS_LIST; + } + } + + record GetPrompt( + Document id, + String name, + Map arguments, + McpMetadata metadata) implements McpCall { + public GetPrompt { + Objects.requireNonNull(name, "name"); + arguments = arguments == null ? Map.of() : Map.copyOf(arguments); + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.PROMPTS_GET; + } + } + + record Complete( + Document id, + CompletionReference reference, + CompletionArgument argument, + McpMetadata metadata) implements McpCall { + public Complete { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.COMPLETION_COMPLETE; + } + } + + record CompletionReference(String type, String name) {} + + record CompletionArgument(String name, String value) {} + + record SetLogLevel(Document id, String level, McpMetadata metadata) implements McpCall { + public SetLogLevel { + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.LOGGING_SET_LEVEL; + } + } + + record ReadResource(Document id, String uri, McpMetadata metadata) implements McpCall { + public ReadResource { + Objects.requireNonNull(uri, "uri"); + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return McpMethod.Standard.RESOURCES_READ; + } + } + + record Notification(McpMethod.Standard method, Document params, McpMetadata metadata) implements McpCall { + public Notification { + Objects.requireNonNull(method, "method"); + params = params == null ? Document.of(Map.of()) : params; + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + if (!method.wireName().startsWith("notifications/")) { + throw new IllegalArgumentException("Not a notification method: " + method.wireName()); + } + } + + @Override + public Document id() { + return null; + } + } + + record ExtensionCall

( + Document id, + McpExtensionMethod

extension, + P parameters, + McpMetadata metadata) implements McpCall { + public ExtensionCall { + Objects.requireNonNull(extension, "extension"); + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + + @Override + public McpMethod method() { + return new McpMethod.Extension(extension.method()); + } + } + + record UnknownCall( + Document id, + McpMethod.Unknown method, + Document params, + McpMetadata metadata) implements McpCall { + public UnknownCall { + Objects.requireNonNull(method, "method"); + params = params == null ? Document.of(Map.of()) : params; + metadata = metadata == null ? McpMetadata.EMPTY : metadata; + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCatalog.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCatalog.java new file mode 100644 index 0000000000..2837044ee3 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCatalog.java @@ -0,0 +1,453 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +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.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.UnaryOperator; +import software.amazon.smithy.java.core.schema.SchemaIndex; +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.PromptInfo; +import software.amazon.smithy.java.mcp.model.ToolInfo; +import software.amazon.smithy.java.server.Service; + +/** + * Thread-safe catalog of local and remote MCP tools and prompts. + * + *

Readers consume immutable snapshots. Mutations rebuild and atomically publish a + * new snapshot so request execution never observes a partially refreshed catalog. + */ +final class McpCatalog implements McpSources { + private static final InternalLogger LOG = InternalLogger.getLogger(McpCatalog.class); + + private final AtomicReference state; + private final AtomicReference> remoteStart = new AtomicReference<>(); + private final AtomicReference> remoteCatalogLoad = new AtomicReference<>(); + private final AtomicReference> remoteInitialization = new AtomicReference<>(); + private final ExecutorService notificationRefreshes = Executors.newVirtualThreadPerTaskExecutor(); + private final Map refreshStates = new ConcurrentHashMap<>(); + private volatile Consumer notificationWriter; + private volatile Consumer responseWriter = ignored -> {}; + private volatile JsonRpcRequest initializeRequest; + private volatile McpProtocol initializeProtocol; + + McpCatalog(Map services, List remoteClients) { + var clients = new HashMap(); + for (var client : remoteClients) { + if (clients.put(client.name(), client) != null) { + throw new IllegalArgumentException("Duplicate remote MCP client: " + client.name()); + } + } + var immutableServices = Map.copyOf(services); + state = new AtomicReference<>(new CatalogState( + immutableServices, + clients, + createLocalSnapshot(immutableServices, Map.of(), Map.of()))); + } + + @Override + public McpSourceSnapshot snapshot() { + return state.get().snapshot(); + } + + @Override + public McpToolDescriptor tool(String name) { + return state.get().snapshot().tools().get(name); + } + + @Override + public McpPromptDescriptor prompt(String normalizedName) { + return state.get().snapshot().prompts().get(normalizedName); + } + + @Override + public Map remoteClients() { + return state.get().remoteClients(); + } + + @Override + public boolean containsServer(String id) { + var current = state.get(); + return current.services().containsKey(id) || current.remoteClients().containsKey(id); + } + + @Override + public void bindTransport( + Consumer notificationWriter, + Consumer responseWriter + ) { + this.notificationWriter = notificationWriter; + this.responseWriter = responseWriter; + runOnce(remoteStart, () -> forEachRemoteInParallel("start", McpRemoteClient::start)); + } + + @Override + public void initializeRemoteClients( + JsonRpcRequest request, + McpProtocol protocol + ) { + initializeRequest = request; + initializeProtocol = protocol; + runOnce(remoteInitialization, () -> { + forEachRemoteInParallel( + "initialize", + client -> initializeAndRefresh(client, request, protocol, responseWriter)); + remoteCatalogLoad.compareAndSet(null, CompletableFuture.completedFuture(null)); + }); + } + + @Override + public void ensureRemoteCatalogLoaded() { + runOnce(remoteCatalogLoad, () -> forEachRemoteInParallel("refresh", this::refresh)); + } + + @Override + public void addService(String id, Service service) { + updateState(current -> { + var services = new HashMap<>(current.services()); + services.put(id, service); + + var schemaIndex = createSchemaIndex(services); + var schemaFactory = new McpSchemaFactory(schemaIndex); + + var tools = new HashMap<>(current.snapshot().tools()); + tools.entrySet() + .removeIf(entry -> entry.getValue().serverId().equals(id) + && entry.getValue().target() instanceof McpToolDescriptor.LocalTarget); + tools.putAll(schemaFactory.createTools(Map.of(id, service))); + tools.putAll(remoteTools(current.snapshot().tools())); + + return new CatalogState( + services, + current.remoteClients(), + new McpSourceSnapshot( + Map.copyOf(tools), + createPromptSnapshot(services, remotePrompts(current.snapshot().prompts())), + new SmithyDocumentAdapter(schemaIndex))); + }); + } + + @Override + public void addRemoteClient(McpRemoteClient client) { + updateState(current -> { + var clients = new HashMap<>(current.remoteClients()); + if (clients.put(client.name(), client) != null) { + throw new IllegalArgumentException("Duplicate remote MCP client: " + client.name()); + } + return new CatalogState(current.services(), clients, current.snapshot()); + }); + + try { + client.start(); + var currentInitializeRequest = initializeRequest; + if (currentInitializeRequest != null) { + initializeAndRefresh( + client, + currentInitializeRequest, + initializeProtocol, + responseWriter); + } else { + refresh(client); + } + } catch (RuntimeException e) { + LOG.error("Failed to add remote MCP client: " + client.name(), e); + } + } + + @Override + public Map headerParameters(String toolName) { + var tool = state.get().snapshot().tools().get(toolName); + return tool == null ? Map.of() : tool.headerParameters(); + } + + @Override + public void close() { + notificationRefreshes.shutdownNow(); + remoteClients().values().forEach(client -> { + try { + client.close(); + } catch (RuntimeException e) { + LOG.error("Failed to close remote MCP client: " + client.name(), e); + } + }); + } + + private void initializeAndRefresh( + McpRemoteClient client, + JsonRpcRequest request, + McpProtocol protocol, + Consumer responseWriter + ) { + client.initialize( + responseWriter, + notification -> onRemoteNotification(client, notification), + request, + protocol); + refresh(client); + } + + private void onRemoteNotification(McpRemoteClient client, JsonRpcRequest notification) { + if (McpMethod.Standard.NOTIFICATIONS_TOOLS_LIST_CHANGED.wireName().equals(notification.getMethod())) { + scheduleRefresh(client); + } + var writer = notificationWriter; + if (writer != null) { + writer.accept(notification); + } + } + + private void scheduleRefresh(McpRemoteClient client) { + var state = refreshStates.computeIfAbsent(client, ignored -> new RefreshState()); + if (state.request()) { + notificationRefreshes.submit(() -> runScheduledRefreshes(client, state)); + } + } + + private void runScheduledRefreshes(McpRemoteClient client, RefreshState state) { + while (state.takeRequest()) { + refresh(client); + } + } + + private void refresh(McpRemoteClient client) { + List remoteTools = List.of(); + boolean toolsLoaded = false; + try { + remoteTools = List.copyOf(client.listTools()); + toolsLoaded = true; + } catch (RuntimeException e) { + LOG.error("Failed to refresh tools from remote MCP client: " + client.name(), e); + } + + List remotePrompts = List.of(); + boolean promptsLoaded = false; + try { + remotePrompts = List.copyOf(client.listPrompts()); + promptsLoaded = true; + } catch (RuntimeException e) { + LOG.error("Failed to refresh prompts from remote MCP client: " + client.name(), e); + } + + mergeRemoteSnapshot(client, remoteTools, toolsLoaded, remotePrompts, promptsLoaded); + } + + private void mergeRemoteSnapshot( + McpRemoteClient client, + List remoteTools, + boolean toolsLoaded, + List remotePrompts, + boolean promptsLoaded + ) { + updateState(current -> { + var tools = new HashMap<>(current.snapshot().tools()); + if (toolsLoaded) { + tools.entrySet() + .removeIf(entry -> entry.getValue().target() instanceof McpToolDescriptor.RemoteTarget remote + && remote.client() == client); + for (var info : remoteTools) { + tools.put( + info.getName(), + new McpToolDescriptor( + info, + client.name(), + new McpToolDescriptor.RemoteTarget(client), + McpHttpBinding.headerParameters(info))); + } + } + + var prompts = new HashMap<>(current.snapshot().prompts()); + if (promptsLoaded) { + prompts.entrySet().removeIf(entry -> entry.getValue().remoteClient() == client); + for (var info : remotePrompts) { + prompts.putIfAbsent( + PromptLoader.normalize(info.getName()), + new McpPromptDescriptor(new Prompt(info, client), client)); + } + } + + return new CatalogState( + current.services(), + current.remoteClients(), + new McpSourceSnapshot( + Map.copyOf(tools), + Map.copyOf(prompts), + current.snapshot().documentAdapter())); + }); + } + + private void forEachRemoteInParallel( + String action, + Consumer operation + ) { + var clients = remoteClients().values(); + try (var executor = Executors.newVirtualThreadPerTaskExecutor()) { + var tasks = clients.stream() + .map(client -> executor.submit(() -> { + try { + operation.accept(client); + } catch (RuntimeException e) { + LOG.error("Failed to " + action + " remote MCP client: " + client.name(), e); + } + })) + .toList(); + for (var task : tasks) { + try { + task.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("Interrupted while waiting for remote MCP clients", e); + } catch (ExecutionException e) { + throw new McpRemoteException("Unexpected remote MCP task failure", e.getCause()); + } + } + } + } + + private void runOnce( + AtomicReference> state, + Runnable operation + ) { + var created = new CompletableFuture(); + var active = state.compareAndExchange(null, created); + if (active != null) { + await(active); + return; + } + + try { + operation.run(); + created.complete(null); + } catch (RuntimeException e) { + created.completeExceptionally(e); + state.compareAndSet(created, null); + throw e; + } + } + + private void await(CompletableFuture operation) { + try { + operation.join(); + } catch (CompletionException e) { + if (e.getCause() instanceof RuntimeException runtimeException) { + throw runtimeException; + } + throw e; + } + } + + private void updateState(UnaryOperator update) { + while (true) { + var current = state.get(); + var updated = update.apply(current); + if (state.compareAndSet(current, updated)) { + return; + } + } + } + + private McpSourceSnapshot createLocalSnapshot( + Map services, + Map remoteTools, + Map remotePrompts + ) { + var schemaIndex = createSchemaIndex(services); + var schemaFactory = new McpSchemaFactory(schemaIndex); + var tools = new HashMap<>(schemaFactory.createTools(services)); + tools.putAll(remoteTools); + + return new McpSourceSnapshot( + Map.copyOf(tools), + createPromptSnapshot(services, remotePrompts), + new SmithyDocumentAdapter(schemaIndex)); + } + + private SchemaIndex createSchemaIndex(Map services) { + return SchemaIndex.compose( + services.values().stream().map(Service::schemaIndex).toArray(SchemaIndex[]::new)); + } + + private Map createPromptSnapshot( + Map services, + Map remotePrompts + ) { + var prompts = new HashMap(); + for (var entry : PromptLoader.loadPrompts(services.values()).entrySet()) { + prompts.put(entry.getKey(), new McpPromptDescriptor(entry.getValue(), null)); + } + remotePrompts.forEach(prompts::putIfAbsent); + return Map.copyOf(prompts); + } + + private Map remoteTools(Map tools) { + var result = new HashMap(); + tools.forEach((name, tool) -> { + if (tool.target() instanceof McpToolDescriptor.RemoteTarget) { + result.put(name, tool); + } + }); + return result; + } + + private Map remotePrompts(Map prompts) { + var result = new HashMap(); + prompts.forEach((name, prompt) -> { + if (prompt.remoteClient() != null) { + result.put(name, prompt); + } + }); + return result; + } + + private record CatalogState( + Map services, + Map remoteClients, + McpSourceSnapshot snapshot) { + private CatalogState { + services = Map.copyOf(services); + remoteClients = Map.copyOf(remoteClients); + } + } + + private static final class RefreshState { + private static final int RUNNING = 1; + private static final int REQUESTED = 1 << 1; + + private final AtomicInteger state = new AtomicInteger(); + + boolean request() { + while (true) { + var current = state.get(); + var updated = current | RUNNING | REQUESTED; + if (state.compareAndSet(current, updated)) { + return (current & RUNNING) == 0; + } + } + } + + boolean takeRequest() { + while (true) { + var current = state.get(); + var hasRequest = (current & REQUESTED) != 0; + var updated = hasRequest ? current & ~REQUESTED : 0; + if (state.compareAndSet(current, updated)) { + return hasRequest; + } + } + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpDomainOperations.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpDomainOperations.java new file mode 100644 index 0000000000..e64a9d001a --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpDomainOperations.java @@ -0,0 +1,221 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.List; +import java.util.Map; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.Capabilities; +import software.amazon.smithy.java.mcp.model.InitializeResult; +import software.amazon.smithy.java.mcp.model.ListPromptsResult; +import software.amazon.smithy.java.mcp.model.ListToolsResult; +import software.amazon.smithy.java.mcp.model.Prompts; +import software.amazon.smithy.java.mcp.model.ServerInfo; +import software.amazon.smithy.java.mcp.model.ToolInfo; +import software.amazon.smithy.java.mcp.model.Tools; + +/** + * Protocol-independent implementation of MCP domain operations. + */ +final class McpDomainOperations implements McpOperations { + private final McpSources sources; + private final McpWireCodec wireCodec; + private final McpServerIdentity identity; + private final ToolFilter toolFilter; + private final McpMetricsObserver metricsObserver; + private final McpToolExecutor toolExecutor; + private final McpProtocolRegistry protocols; + + McpDomainOperations( + McpSources sources, + McpWireCodec wireCodec, + McpServerIdentity identity, + ToolFilter toolFilter, + McpMetricsObserver metricsObserver, + McpInterceptor interceptor, + McpProtocolRegistry protocols + ) { + this.sources = sources; + this.wireCodec = wireCodec; + this.identity = identity; + this.toolFilter = toolFilter; + this.metricsObserver = metricsObserver; + this.protocols = protocols; + this.toolExecutor = new McpToolExecutor(sources, wireCodec, interceptor, protocols); + } + + @Override + public McpOutcome initialize(McpCall.Initialize call, McpRequestContext context) { + observeInitialize(call); + sources.initializeRemoteClients( + wireCodec.encode(call), + protocols.require(context.protocolVersion())); + + var result = InitializeResult.builder() + .protocolVersion(context.protocolVersion().identifier()) + .capabilities(Capabilities.builder() + .completions(Document.of(Map.of())) + .logging(Document.of(Map.of())) + .tools(Tools.builder().listChanged(true).build()) + .prompts(Prompts.builder().listChanged(true).build()) + .build()) + .serverInfo(ServerInfo.builder() + .name(identity.name()) + .version(identity.version()) + .build()) + .build(); + return new McpOutcome.Success(call.id(), Document.of(result)); + } + + @Override + public McpOutcome ping(McpCall.Ping call, McpRequestContext context) { + return new McpOutcome.Success(call.id(), Document.of(Map.of())); + } + + @Override + public McpOutcome discover(McpCall.Discover call, McpRequestContext context) { + sources.ensureRemoteCatalogLoaded(); + var capabilities = Document.of(Map.of( + "completions", + Document.of(Map.of()), + "tools", + Document.of(Map.of()), + "prompts", + Document.of(Map.of()))); + return new McpOutcome.Success( + call.id(), + Document.of(Map.of( + "supportedVersions", + Document.of(protocols.supportedIdentifiers() + .stream() + .map(Document::of) + .toList()), + "capabilities", + capabilities))); + } + + @Override + public McpOutcome listTools(McpCall.ListTools call, McpRequestContext context) { + sources.ensureRemoteCatalogLoaded(); + var protocol = protocols.require(context.protocolVersion()); + var tools = sources.snapshot() + .tools() + .values() + .stream() + .filter(tool -> toolFilter.allowTool(tool.serverId(), tool.info().getName())) + .map(tool -> projectTool(tool.info(), protocol)) + .toList(); + return new McpOutcome.Success( + call.id(), + Document.of(ListToolsResult.builder().tools(tools).build())); + } + + @Override + public McpOutcome callTool(McpCall.CallTool call, McpRequestContext context) { + sources.ensureRemoteCatalogLoaded(); + if (metricsObserver != null) { + metricsObserver.onToolCall(call.method().wireName(), call.name()); + } + return toolExecutor.execute(call, context); + } + + @Override + public McpOutcome listPrompts(McpCall.ListPrompts call, McpRequestContext context) { + sources.ensureRemoteCatalogLoaded(); + var prompts = sources.snapshot() + .prompts() + .values() + .stream() + .map(descriptor -> descriptor.prompt().promptInfo()) + .toList(); + return new McpOutcome.Success( + call.id(), + Document.of(ListPromptsResult.builder().prompts(prompts).build())); + } + + @Override + public McpOutcome getPrompt(McpCall.GetPrompt call, McpRequestContext context) { + sources.ensureRemoteCatalogLoaded(); + var prompt = sources.prompt(PromptLoader.normalize(call.name())); + if (prompt == null) { + return new McpOutcome.Failure( + call.id(), + new McpError(-32602, "Prompt not found: " + call.name(), null)); + } + var arguments = call.arguments().isEmpty() ? null : Document.of(call.arguments()); + return new McpOutcome.Success( + call.id(), + Document.of(prompt.prompt().getPromptResult(arguments, call.id()))); + } + + @Override + public McpOutcome complete(McpCall.Complete call, McpRequestContext context) { + var completion = Document.of(Map.of( + "values", + Document.of(List.of()), + "total", + Document.of(0), + "hasMore", + Document.of(false))); + return new McpOutcome.Success(call.id(), Document.of(Map.of("completion", completion))); + } + + @Override + public McpOutcome setLogLevel(McpCall.SetLogLevel call, McpRequestContext context) { + return new McpOutcome.Success(call.id(), Document.of(Map.of())); + } + + @Override + public McpOutcome readResource(McpCall.ReadResource call, McpRequestContext context) { + throw new UnsupportedOperationException("resources/read is not implemented"); + } + + private ToolInfo projectTool(ToolInfo tool, McpProtocol protocol) { + boolean stripOutput = !protocol.supportsOutputSchema() && tool.getOutputSchema() != null; + boolean stripAnnotations = !protocol.supportsAnnotations() && tool.getAnnotations() != null; + if (!stripOutput && !stripAnnotations) { + return tool; + } + var builder = tool.toBuilder(); + if (stripOutput) { + builder.outputSchema(null); + } + if (stripAnnotations) { + builder.annotations(null); + } + return builder.build(); + } + + private void observeInitialize(McpCall.Initialize call) { + if (metricsObserver == null) { + return; + } + var capabilities = call.capabilities(); + var clientInfo = call.clientInfo(); + boolean rootsListChanged = capabilities != null + && capabilities.getMember("roots") != null + && capabilities.getMember("roots").getMember("listChanged") != null + && capabilities.getMember("roots").getMember("listChanged").asBoolean(); + boolean sampling = capabilities != null && capabilities.getMember("sampling") != null; + boolean elicitation = capabilities != null && capabilities.getMember("elicitation") != null; + metricsObserver.onInitialize( + call.method().wireName(), + call.requestedVersion().identifier(), + rootsListChanged, + sampling, + elicitation, + stringMember(clientInfo, "name"), + stringMember(clientInfo, "title")); + } + + private String stringMember(Document document, String name) { + if (document == null) { + return null; + } + var member = document.getMember(name); + return member == null ? null : member.asString(); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpEngine.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpEngine.java new file mode 100644 index 0000000000..b588db5aa0 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpEngine.java @@ -0,0 +1,335 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.function.Consumer; +import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.serde.document.Document; +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.server.Service; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Blocking, transport-independent MCP execution engine. + * + *

The engine operates on typed calls and outcomes. Transports own concurrency and + * should invoke this blocking API from virtual threads when concurrent execution is + * desired. + */ +@SmithyUnstableApi +public final class McpEngine implements AutoCloseable { + private static final InternalLogger LOG = InternalLogger.getLogger(McpEngine.class); + + private final McpSources sources; + private final McpDomainOperations operations; + private final McpWireCodec wireCodec; + private final McpInterceptor interceptor; + private final McpServerIdentity identity; + private final McpProtocolRegistry protocols; + + private McpEngine(Builder builder) { + identity = new McpServerIdentity(builder.name, builder.version); + protocols = McpProtocolRegistry.create( + builder.protocols.values(), + builder.protocolOverrides.values(), + builder.discoverProtocols); + wireCodec = new McpWireCodec(builder.extensions); + interceptor = builder.interceptor; + sources = new McpCatalog(builder.services, builder.remoteClients); + operations = new McpDomainOperations( + sources, + wireCodec, + identity, + builder.toolFilter, + builder.metricsObserver, + interceptor, + protocols); + } + + /** + * Executes a typed call synchronously. + */ + public McpOutcome execute(McpCall call, McpRequestContext requestContext) { + var executionContext = new McpExecutionContext(call, requestContext); + McpOutcome outcome = null; + RuntimeException error = null; + + try { + interceptor.readBeforeExecution(executionContext); + call = interceptor.modifyBeforeExecution(executionContext); + executionContext = executionContext.withCall(call); + + var protocol = protocols.require(requestContext.protocolVersion()); + protocol.validate(call, requestContext); + outcome = protocol.dispatch(call, operations, requestContext); + if (outcome instanceof McpOutcome.Success(Document id, Document result)) { + outcome = new McpOutcome.Success( + id, + protocol.decorateResult(result, call.method(), identity)); + } + } catch (RuntimeException e) { + error = e; + } + + try { + interceptor.readAfterExecution(executionContext, outcome, error); + } catch (RuntimeException e) { + error = preserveOriginal(error, e); + } + + try { + return interceptor.modifyAfterExecution(executionContext, outcome, error); + } catch (RuntimeException e) { + return errorOutcome(call, e); + } + } + + /** + * Executes a decoded JSON-RPC request with an explicit protocol claim. + * + *

This is primarily useful for transport adapters. Application code should + * prefer the typed-call overload. + */ + public JsonRpcResponse execute(JsonRpcRequest request, ProtocolVersion protocolVersion) { + var session = newSession(); + var outcome = execute(request, session, protocolVersion, McpTransportContext.STDIO); + return encode(outcome); + } + + McpOutcome execute( + JsonRpcRequest request, + McpSession session, + ProtocolVersion transportClaim, + McpTransportContext transportContext + ) { + final McpCall call; + try { + call = wireCodec.decode(request); + } catch (RuntimeException e) { + return errorOutcome(request.getId(), e); + } + + final ProtocolVersion version; + try { + version = session.negotiate(call, transportClaim); + } catch (RuntimeException e) { + return errorOutcome(call, e); + } + return execute(call, new McpRequestContext(version, transportContext, Context.create())); + } + + JsonRpcRequest encode(McpCall call) { + return wireCodec.encode(call); + } + + JsonRpcResponse encode(McpOutcome outcome) { + return wireCodec.encode(outcome); + } + + McpOutcome decode(JsonRpcResponse response) { + return wireCodec.decode(response); + } + + McpServerIdentity identity() { + return identity; + } + + McpSession newSession() { + return new McpSession(protocols); + } + + McpProtocol protocol(ProtocolVersion version) { + return protocols.require(version); + } + + McpProtocol findProtocol(ProtocolVersion version) { + return protocols.find(version); + } + + void bindTransport( + Consumer notificationWriter, + Consumer responseWriter + ) { + sources.bindTransport(notificationWriter, responseWriter); + } + + void addService(String id, Service service) { + sources.addService(id, service); + } + + void addRemoteClient(McpRemoteClient client) { + sources.addRemoteClient(client); + } + + boolean containsServer(String id) { + return sources.containsServer(id); + } + + Map remoteClients() { + return sources.remoteClients(); + } + + Map headerParameters(String toolName) { + return sources.headerParameters(toolName); + } + + @Override + public void close() { + sources.close(); + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private final Map services = new HashMap<>(); + private final List remoteClients = new ArrayList<>(); + private final Map> extensions = new HashMap<>(); + private final Map protocols = new LinkedHashMap<>(); + private final Map protocolOverrides = new LinkedHashMap<>(); + private McpInterceptor interceptor = McpInterceptor.NOOP; + private String name = "mcp-server"; + private String version = "1.0.0"; + private ToolFilter toolFilter = (serverId, toolName) -> true; + private McpMetricsObserver metricsObserver; + private boolean discoverProtocols = true; + + public Builder services(Map services) { + this.services.clear(); + this.services.putAll(services); + return this; + } + + public Builder addService(String id, Service service) { + services.put(id, service); + return this; + } + + public Builder remoteClients(List remoteClients) { + this.remoteClients.clear(); + this.remoteClients.addAll(remoteClients); + return this; + } + + public Builder addRemoteClient(McpRemoteClient remoteClient) { + remoteClients.add(remoteClient); + return this; + } + + public Builder addExtension(McpExtensionMethod extension) { + Objects.requireNonNull(extension, "extension"); + if (McpMethod.parse(extension.method()) instanceof McpMethod.Standard) { + throw new IllegalArgumentException("Cannot replace standard MCP method: " + extension.method()); + } + if (extensions.put(extension.method(), extension) != null) { + throw new IllegalArgumentException("Duplicate MCP extension method: " + extension.method()); + } + return this; + } + + public Builder addProtocol(ExtensionMcpProtocol protocol) { + putProtocol(protocols, protocol, "protocol"); + return this; + } + + public Builder overrideProtocol(ExtensionMcpProtocol protocol) { + putProtocol(protocolOverrides, protocol, "protocol override"); + return this; + } + + public Builder discoverProtocols(boolean discoverProtocols) { + this.discoverProtocols = discoverProtocols; + return this; + } + + public Builder name(String name) { + this.name = Objects.requireNonNull(name, "name"); + return this; + } + + public Builder version(String version) { + this.version = Objects.requireNonNull(version, "version"); + return this; + } + + public Builder toolFilter(ToolFilter toolFilter) { + this.toolFilter = Objects.requireNonNull(toolFilter, "toolFilter"); + return this; + } + + public Builder metricsObserver(McpMetricsObserver metricsObserver) { + this.metricsObserver = metricsObserver; + return this; + } + + public Builder interceptor(McpInterceptor interceptor) { + this.interceptor = Objects.requireNonNull(interceptor, "interceptor"); + return this; + } + + public McpEngine build() { + return new McpEngine(this); + } + + private void putProtocol( + Map destination, + ExtensionMcpProtocol protocol, + String kind + ) { + Objects.requireNonNull(protocol, kind); + Objects.requireNonNull(protocol.id(), kind + " id"); + if (destination.put(protocol.id(), protocol) != null) { + throw new IllegalArgumentException( + "Duplicate MCP " + kind + ": " + protocol.id().identifier()); + } + } + } + + private McpOutcome errorOutcome(McpCall call, RuntimeException error) { + if (call.id() == null) { + return McpOutcome.NoResponse.INSTANCE; + } + if (error instanceof McpUnsupportedMethodException) { + return new McpOutcome.Failure( + call.id(), + new McpError(-32601, "Method not found: " + call.method().wireName(), null)); + } + return errorOutcome(call.id(), error); + } + + private McpOutcome errorOutcome(Document id, RuntimeException error) { + if (error instanceof McpProtocolException protocolError) { + return new McpOutcome.Failure( + id, + new McpError(protocolError.code(), protocolError.getMessage(), protocolError.data())); + } + LOG.error("Unexpected MCP engine error", error); + return new McpOutcome.Failure(id, new McpError(-32603, "Internal error", null)); + } + + private RuntimeException preserveOriginal( + RuntimeException original, + RuntimeException afterHookFailure + ) { + if (original == null) { + return afterHookFailure; + } + if (original != afterHookFailure) { + original.addSuppressed(afterHookFailure); + } + return original; + } + +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpError.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpError.java new file mode 100644 index 0000000000..9c5f003dcb --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpError.java @@ -0,0 +1,20 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * A semantic MCP error independent of a transport encoding. + */ +@SmithyUnstableApi +public record McpError(int code, String message, Document data) { + public McpError { + Objects.requireNonNull(message, "message"); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionContext.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionContext.java new file mode 100644 index 0000000000..aeb4fdda14 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionContext.java @@ -0,0 +1,24 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Typed data exposed to execution interceptors. + */ +@SmithyUnstableApi +public record McpExecutionContext(McpCall call, McpRequestContext requestContext) { + public McpExecutionContext { + Objects.requireNonNull(call, "call"); + Objects.requireNonNull(requestContext, "requestContext"); + } + + McpExecutionContext withCall(McpCall call) { + return this.call == call ? this : new McpExecutionContext(call, requestContext); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionHook.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionHook.java deleted file mode 100644 index 7a417b4d00..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionHook.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import software.amazon.smithy.java.context.Context; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; -import software.amazon.smithy.utils.SmithyUnstableApi; - -/** - * Hook data available at the execution level. Passed to execution-scoped hooks in - * {@link McpServerInterceptor}. - * - *

The {@link #context()} provides a per-request key-value store for passing state - * between hooks. For example, a telemetry interceptor can stash a start timestamp in - * {@code readBeforeExecution} and retrieve it in {@code readAfterExecution}. - */ -@SmithyUnstableApi -public class McpExecutionHook { - - private final JsonRpcRequest request; - private final ProtocolVersion protocolVersion; - private final Context context; - - McpExecutionHook(JsonRpcRequest request, ProtocolVersion protocolVersion, Context context) { - this.request = request; - this.protocolVersion = protocolVersion; - this.context = context; - } - - /** - * The JSON-RPC request being handled. - */ - public JsonRpcRequest request() { - return request; - } - - /** - * Returns a new hook with the given request, or the same hook if unchanged. - */ - public McpExecutionHook withRequest(JsonRpcRequest request) { - return this.request == request ? this : new McpExecutionHook(request, protocolVersion, context); - } - - /** - * The MCP protocol version for this request. - */ - public ProtocolVersion protocolVersion() { - return protocolVersion; - } - - /** - * Per-request context for passing state between hooks. - */ - public Context context() { - return context; - } -} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExtensionMethod.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExtensionMethod.java new file mode 100644 index 0000000000..997dd90a20 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExtensionMethod.java @@ -0,0 +1,30 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Defines a typed custom MCP method without opening the built-in call hierarchy. + */ +@SmithyUnstableApi +public interface McpExtensionMethod

{ + String method(); + + P decode(Document params); + + /** + * Encodes typed parameters when an extension call is sent to another peer. + * + *

Inbound-only extensions do not need to override this method. + */ + default Document encode(P params) { + throw new UnsupportedOperationException("Extension does not support outbound encoding: " + method()); + } + + McpOutcome execute(McpCall.ExtensionCall

call, McpRequestContext context); +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpBinding.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpBinding.java new file mode 100644 index 0000000000..efc3a9ab83 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpBinding.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 java.nio.charset.StandardCharsets; +import java.util.Base64; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.http.api.HeaderName; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.java.mcp.model.JsonRpcResponse; +import software.amazon.smithy.java.mcp.model.ToolInfo; +import software.amazon.smithy.model.shapes.ShapeType; + +/** + * Shared Streamable HTTP binding rules used by both HTTP peers. + */ +final class McpHttpBinding { + static final HeaderName PROTOCOL_VERSION = HeaderName.of("mcp-protocol-version"); + static final HeaderName SESSION_ID = HeaderName.of("mcp-session-id"); + static final HeaderName METHOD = HeaderName.of("mcp-method"); + static final HeaderName NAME = HeaderName.of("mcp-name"); + + private McpHttpBinding() {} + + static boolean usesMethodHeaders(McpProtocol protocol) { + return protocol.usesHttpMethodHeaders(); + } + + static boolean isInitialize(JsonRpcRequest request) { + return McpMethod.parse(request.getMethod()) == McpMethod.Standard.INITIALIZE; + } + + static boolean isToolCall(JsonRpcRequest request) { + return McpMethod.parse(request.getMethod()) == McpMethod.Standard.TOOLS_CALL; + } + + static String requestName(JsonRpcRequest request) { + var params = request.getParams(); + if (params == null) { + return null; + } + return switch (McpMethod.parse(request.getMethod())) { + case McpMethod.Standard.TOOLS_CALL, McpMethod.Standard.PROMPTS_GET -> + stringMember(params, "name"); + case McpMethod.Standard.RESOURCES_READ -> stringMember(params, "uri"); + default -> null; + }; + } + + static String protocolVersionFromMetadata(JsonRpcRequest request) { + var params = request.getParams(); + var metadata = params == null ? null : params.getMember("_meta"); + return metadata == null ? null : stringMember(metadata, McpWireNames.PROTOCOL_VERSION); + } + + static String stringMember(Document document, String name) { + var member = document.getMember(name); + return member == null || !member.isType(ShapeType.STRING) ? null : member.asString(); + } + + static String firstHeader(Map> headers, String name) { + var values = headers.get(name.toLowerCase(Locale.ROOT)); + return values == null || values.isEmpty() ? null : values.getFirst(); + } + + static Map> normalizeHeaders(Map> headers) { + var normalized = new HashMap>(); + headers.forEach((name, values) -> normalized.put(name.toLowerCase(Locale.ROOT), List.copyOf(values))); + return Map.copyOf(normalized); + } + + static String encodeParameter(String value) { + if (value.startsWith("=?base64?")) { + return encodeBase64(value); + } + for (int index = 0; index < value.length(); index++) { + var character = value.charAt(index); + if (character < 0x20 || character > 0x7e) { + return encodeBase64(value); + } + } + return value; + } + + private static String encodeBase64(String value) { + return "=?base64?" + + Base64.getEncoder().encodeToString(value.getBytes(StandardCharsets.UTF_8)) + + "?="; + } + + static String decodeParameter(String value) { + if (!value.startsWith("=?base64?") || !value.endsWith("?=")) { + return value; + } + + var encoded = value.substring("=?base64?".length(), value.length() - 2); + if (encoded.length() % 4 != 0 + || !encoded.matches("(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?")) { + throw new IllegalArgumentException("Invalid Base64"); + } + return new String(Base64.getDecoder().decode(encoded), StandardCharsets.UTF_8); + } + + static Map headerParameters(ToolInfo tool) { + var result = new HashMap(); + var inputSchema = tool.getInputSchema(); + if (inputSchema == null || inputSchema.getProperties() == null) { + return result; + } + + for (var entry : inputSchema.getProperties().entrySet()) { + var suffix = entry.getValue().getMember("x-mcp-header"); + if (suffix != null + && suffix.isType(ShapeType.STRING) + && suffix.asString().matches("[A-Za-z0-9][A-Za-z0-9_-]*")) { + result.put(entry.getKey(), suffix.asString()); + } + } + return Map.copyOf(result); + } + + static int statusCode( + JsonRpcResponse response, + boolean statelessClaim, + boolean protocolVersionHeaderPresent + ) { + if (response.getError() == null) { + return 200; + } + if (protocolVersionHeaderPresent && response.getError().getCode() == -32022) { + return 400; + } + if (!statelessClaim) { + return 200; + } + return switch (response.getError().getCode()) { + case -32601 -> 404; + case -32602, -32020, -32021, -32022 -> 400; + default -> 200; + }; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpHandler.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpHandler.java new file mode 100644 index 0000000000..40c4a59766 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpHandler.java @@ -0,0 +1,266 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.List; +import java.util.Map; +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.model.shapes.ShapeType; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Maps Streamable HTTP request metadata to the transport-independent MCP service. + */ +@SmithyUnstableApi +public final class McpHttpHandler { + private static final int HEADER_MISMATCH_ERROR_CODE = -32020; + private final McpEngine engine; + private final boolean loopbackOnly; + + public McpHttpHandler(McpEngine engine) { + this(engine, false); + } + + private McpHttpHandler(McpEngine engine, boolean loopbackOnly) { + this.engine = engine; + this.loopbackOnly = loopbackOnly; + engine.bindTransport(ignored -> {}, ignored -> {}); + } + + /** + * Creates a handler for a server bound exclusively to a loopback interface. + * + *

The handler rejects non-loopback Host and Origin headers to protect local + * MCP servers from DNS rebinding attacks. + */ + public static McpHttpHandler forLoopback(McpEngine engine) { + return new McpHttpHandler(engine, true); + } + + /** + * Handles a decoded Streamable HTTP request. + * + * @param request JSON-RPC request body. + * @param headers HTTP request headers. + * @return HTTP status and optional JSON-RPC response body. + */ + public Response handle( + JsonRpcRequest request, + Map> headers + ) { + headers = McpHttpBinding.normalizeHeaders(headers); + var bodyVersion = McpHttpBinding.protocolVersionFromMetadata(request); + var headerVersion = McpHttpBinding.firstHeader(headers, "mcp-protocol-version"); + var protocolVersion = resolveProtocolVersion(request, bodyVersion, headerVersion); + var modernClaim = isModernClaim(bodyVersion, headerVersion); + + if (loopbackOnly) { + var hostError = validateLoopbackHeaders(request, headers); + if (hostError != null) { + return new Response(400, hostError); + } + } + + if (modernClaim) { + var headerError = validateModernHeaders(request, headers, bodyVersion, headerVersion); + if (headerError != null) { + return new Response(400, headerError); + } + + var parameterError = validateMcpParameterHeaders(request, headers); + if (parameterError != null) { + return new Response(400, parameterError); + } + } + + var outcome = engine.execute( + request, + engine.newSession(), + protocolVersion, + new McpTransportContext.Http(headers, loopbackOnly)); + var response = engine.encode(outcome); + if (response == null) { + return new Response(202, null); + } + return new Response( + McpHttpBinding.statusCode(response, modernClaim, headerVersion != null), + response); + } + + private ProtocolVersion resolveProtocolVersion( + JsonRpcRequest request, + String bodyVersion, + String headerVersion + ) { + if (bodyVersion != null) { + return ProtocolVersion.parse(bodyVersion); + } + if (McpHttpBinding.isInitialize(request)) { + var params = request.getParams(); + var initializeVersion = params == null ? null : params.getMember("protocolVersion"); + var identifier = initializeVersion == null + ? null + : McpHttpBinding.stringMember(params, "protocolVersion"); + return identifier == null + ? ProtocolVersion.defaultVersion() + : ProtocolVersion.parse(identifier); + } + return headerVersion == null + ? ProtocolVersion.defaultVersion() + : ProtocolVersion.parse(headerVersion); + } + + private boolean isModernClaim(String bodyVersion, String headerVersion) { + if (bodyVersion != null) { + return true; + } + if (headerVersion == null) { + return false; + } + var protocol = engine.findProtocol(ProtocolVersion.parse(headerVersion)); + return protocol != null && McpHttpBinding.usesMethodHeaders(protocol); + } + + private JsonRpcResponse validateModernHeaders( + JsonRpcRequest request, + Map> headers, + String bodyVersion, + String headerVersion + ) { + if (headerVersion == null) { + return headerMismatch(request, "Missing MCP-Protocol-Version header"); + } + if (bodyVersion != null && !bodyVersion.equals(headerVersion)) { + return headerMismatch(request, "MCP-Protocol-Version header does not match request metadata"); + } + + var method = McpHttpBinding.firstHeader(headers, "mcp-method"); + if (!request.getMethod().equals(method)) { + return headerMismatch(request, "Mcp-Method header does not match the JSON-RPC method"); + } + + var expectedName = McpHttpBinding.requestName(request); + var actualName = McpHttpBinding.firstHeader(headers, "mcp-name"); + if (expectedName != null && !expectedName.equals(actualName)) { + return headerMismatch(request, "Mcp-Name header does not match the request parameters"); + } + if (expectedName == null && actualName != null) { + return headerMismatch(request, "Mcp-Name header is not valid for this method"); + } + return null; + } + + private JsonRpcResponse validateMcpParameterHeaders( + JsonRpcRequest request, + Map> headers + ) { + if (!McpHttpBinding.isToolCall(request)) { + return null; + } + + var params = request.getParams(); + var toolName = params == null ? null : McpHttpBinding.stringMember(params, "name"); + if (toolName == null) { + return null; + } + + var arguments = params.getMember("arguments"); + for (var entry : engine.headerParameters(toolName).entrySet()) { + var parameterName = entry.getKey(); + var headerName = "Mcp-Param-" + entry.getValue(); + var bodyValue = arguments == null ? null : arguments.getMember(parameterName); + var headerValue = McpHttpBinding.firstHeader(headers, headerName); + + if (bodyValue == null && headerValue == null) { + continue; + } + if (bodyValue == null || headerValue == null || !bodyValue.isType(ShapeType.STRING)) { + return headerMismatch(request, headerName + " does not match the JSON body parameter"); + } + + final String decodedHeader; + try { + decodedHeader = McpHttpBinding.decodeParameter(headerValue); + } catch (IllegalArgumentException e) { + return headerMismatch(request, headerName + " contains invalid Base64"); + } + if (!bodyValue.asString().equals(decodedHeader)) { + return headerMismatch(request, headerName + " does not match the JSON body parameter"); + } + } + return null; + } + + private JsonRpcResponse validateLoopbackHeaders( + JsonRpcRequest request, + Map> headers + ) { + var host = McpHttpBinding.firstHeader(headers, "host"); + if (!isLoopbackAuthority(host)) { + return headerMismatch(request, "Host header is not a loopback address"); + } + + var origin = McpHttpBinding.firstHeader(headers, "origin"); + if (origin != null && !isLoopbackOrigin(origin)) { + return headerMismatch(request, "Origin header is not a loopback origin"); + } + return null; + } + + private boolean isLoopbackOrigin(String value) { + try { + return isLoopbackHost(new URI(value).getHost()); + } catch (URISyntaxException e) { + return false; + } + } + + private boolean isLoopbackAuthority(String value) { + if (value == null || value.isBlank()) { + return false; + } + try { + return isLoopbackHost(new URI("http://" + value).getHost()); + } catch (URISyntaxException e) { + return false; + } + } + + private boolean isLoopbackHost(String host) { + if (host == null) { + return false; + } + if (host.startsWith("[") && host.endsWith("]")) { + host = host.substring(1, host.length() - 1); + } + return "localhost".equalsIgnoreCase(host) + || "127.0.0.1".equals(host) + || "::1".equals(host); + } + + private JsonRpcResponse headerMismatch(JsonRpcRequest request, String message) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .error(JsonRpcErrorResponse.builder() + .code(HEADER_MISMATCH_ERROR_CODE) + .message(message) + .build()) + .build(); + } + + /** + * A Streamable HTTP response. + * + * @param statusCode HTTP status code. + * @param body JSON-RPC body, or {@code null} when no response body is required. + */ + public record Response(int statusCode, JsonRpcResponse body) {} +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptor.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptor.java new file mode 100644 index 0000000000..a8be6d56a7 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptor.java @@ -0,0 +1,77 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.List; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Typed interceptor for MCP execution. + * + *

Calls and outcomes are immutable; modifying hooks return replacement values. + */ +@SmithyUnstableApi +public interface McpInterceptor { + McpInterceptor NOOP = new McpInterceptor() {}; + + static McpInterceptor chain(McpInterceptor... interceptors) { + return chain(List.of(interceptors)); + } + + static McpInterceptor chain(List interceptors) { + return switch (interceptors.size()) { + case 0 -> NOOP; + case 1 -> interceptors.getFirst(); + default -> new McpInterceptorChain(List.copyOf(interceptors)); + }; + } + + default void readBeforeExecution(McpExecutionContext context) {} + + default McpCall modifyBeforeExecution(McpExecutionContext context) { + return context.call(); + } + + default void readAfterExecution( + McpExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) {} + + default McpOutcome modifyAfterExecution( + McpExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + if (error != null) { + throw error; + } + return outcome; + } + + default void readBeforeToolCall(McpToolExecutionContext context) {} + + default McpCall.CallTool modifyBeforeToolCall(McpToolExecutionContext context) { + return context.call(); + } + + default void readAfterToolCall( + McpToolExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) {} + + default McpOutcome modifyAfterToolCall( + McpToolExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + if (error != null) { + throw error; + } + return outcome; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptorChain.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptorChain.java new file mode 100644 index 0000000000..54f25513b6 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptorChain.java @@ -0,0 +1,88 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.List; + +final class McpInterceptorChain implements McpInterceptor { + private final List interceptors; + + McpInterceptorChain(List interceptors) { + this.interceptors = interceptors; + } + + @Override + public void readBeforeExecution(McpExecutionContext context) { + interceptors.forEach(interceptor -> interceptor.readBeforeExecution(context)); + } + + @Override + public McpCall modifyBeforeExecution(McpExecutionContext context) { + var call = context.call(); + for (var interceptor : interceptors) { + call = interceptor.modifyBeforeExecution(context.withCall(call)); + } + return call; + } + + @Override + public void readAfterExecution(McpExecutionContext context, McpOutcome outcome, RuntimeException error) { + interceptors.forEach(interceptor -> interceptor.readAfterExecution(context, outcome, error)); + } + + @Override + public McpOutcome modifyAfterExecution( + McpExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + var current = outcome; + var currentError = error; + for (var interceptor : interceptors) { + current = interceptor.modifyAfterExecution(context, current, currentError); + currentError = null; + } + return current; + } + + @Override + public void readBeforeToolCall(McpToolExecutionContext context) { + interceptors.forEach(interceptor -> interceptor.readBeforeToolCall(context)); + } + + @Override + public McpCall.CallTool modifyBeforeToolCall(McpToolExecutionContext context) { + var call = context.call(); + for (var interceptor : interceptors) { + call = interceptor.modifyBeforeToolCall(context.withCall(call)); + } + return call; + } + + @Override + public void readAfterToolCall( + McpToolExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + interceptors.forEach(interceptor -> interceptor.readAfterToolCall(context, outcome, error)); + } + + @Override + public McpOutcome modifyAfterToolCall( + McpToolExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + var current = outcome; + var currentError = error; + for (var interceptor : interceptors) { + current = interceptor.modifyAfterToolCall(context, current, currentError); + currentError = null; + } + return current; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpJson.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpJson.java new file mode 100644 index 0000000000..249fca010f --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpJson.java @@ -0,0 +1,20 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import software.amazon.smithy.java.json.JsonCodec; +import software.amazon.smithy.java.json.JsonSettings; + +final class McpJson { + static final JsonCodec CODEC = JsonCodec.builder() + .settings(JsonSettings.builder() + .serializeTypeInDocuments(false) + .useJsonName(true) + .build()) + .build(); + + private McpJson() {} +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMetadata.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMetadata.java new file mode 100644 index 0000000000..9f1d5dc565 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMetadata.java @@ -0,0 +1,26 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Map; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Metadata shared by typed MCP calls. + */ +@SmithyUnstableApi +public record McpMetadata( + ProtocolVersion protocolVersion, + Document clientInfo, + Document clientCapabilities, + Map extensions) { + public static final McpMetadata EMPTY = new McpMetadata(null, null, null, Map.of()); + + public McpMetadata { + extensions = extensions == null ? Map.of() : Map.copyOf(extensions); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMethod.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMethod.java new file mode 100644 index 0000000000..05902e8152 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMethod.java @@ -0,0 +1,103 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * A typed MCP method name. + */ +@SmithyUnstableApi +public sealed interface McpMethod permits McpMethod.Standard, McpMethod.Extension, McpMethod.Unknown { + + /** + * Returns the wire method name. + */ + String wireName(); + + /** + * Parses a wire method name. + */ + static McpMethod parse(String wireName) { + Objects.requireNonNull(wireName, "wireName"); + var standard = Standard.fromWireName(wireName); + return standard == null ? new Unknown(wireName) : standard; + } + + /** + * Methods defined by the MCP protocol. + */ + enum Standard implements McpMethod { + INITIALIZE("initialize"), + PING("ping"), + SERVER_DISCOVER("server/discover"), + PROMPTS_LIST("prompts/list"), + PROMPTS_GET("prompts/get"), + COMPLETION_COMPLETE("completion/complete"), + LOGGING_SET_LEVEL("logging/setLevel"), + TOOLS_LIST("tools/list"), + TOOLS_CALL("tools/call"), + RESOURCES_READ("resources/read"), + NOTIFICATIONS_INITIALIZED("notifications/initialized"), + NOTIFICATIONS_TOOLS_LIST_CHANGED("notifications/tools/list_changed"); + + private static final Map BY_WIRE_NAME; + + static { + var methods = new HashMap(); + for (var method : values()) { + if (methods.put(method.wireName, method) != null) { + throw new IllegalStateException("Duplicate MCP method: " + method.wireName); + } + } + BY_WIRE_NAME = Collections.unmodifiableMap(methods); + } + + private final String wireName; + + Standard(String wireName) { + this.wireName = wireName; + } + + @Override + public String wireName() { + return wireName; + } + + static Standard fromWireName(String wireName) { + return BY_WIRE_NAME.get(wireName); + } + } + + /** + * A registered extension method. + */ + record Extension(String wireName) implements McpMethod { + public Extension { + requireWireName(wireName); + } + } + + /** + * An unrecognized method received from a peer. + */ + record Unknown(String wireName) implements McpMethod { + public Unknown { + requireWireName(wireName); + } + } + + private static void requireWireName(String wireName) { + Objects.requireNonNull(wireName, "wireName"); + if (wireName.isBlank()) { + throw new IllegalArgumentException("MCP method name must not be blank"); + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOperations.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOperations.java new file mode 100644 index 0000000000..e95f80a6c4 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOperations.java @@ -0,0 +1,37 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Domain operations available to protocol implementations. + * + *

Protocol implementations decide whether an operation exists and how its result is + * projected. Implementations delegate supported operations to this interface. + */ +@SmithyUnstableApi +public interface McpOperations { + McpOutcome initialize(McpCall.Initialize call, McpRequestContext context); + + McpOutcome ping(McpCall.Ping call, McpRequestContext context); + + McpOutcome discover(McpCall.Discover call, McpRequestContext context); + + McpOutcome listTools(McpCall.ListTools call, McpRequestContext context); + + McpOutcome callTool(McpCall.CallTool call, McpRequestContext context); + + McpOutcome listPrompts(McpCall.ListPrompts call, McpRequestContext context); + + McpOutcome getPrompt(McpCall.GetPrompt call, McpRequestContext context); + + McpOutcome complete(McpCall.Complete call, McpRequestContext context); + + McpOutcome setLogLevel(McpCall.SetLogLevel call, McpRequestContext context); + + McpOutcome readResource(McpCall.ReadResource call, McpRequestContext context); +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOutcome.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOutcome.java new file mode 100644 index 0000000000..27bdf91adf --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOutcome.java @@ -0,0 +1,33 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * The result of blocking MCP execution. + */ +@SmithyUnstableApi +public sealed interface McpOutcome permits McpOutcome.Success, McpOutcome.Failure, McpOutcome.NoResponse { + + record Success(Document id, Document result) implements McpOutcome { + public Success { + Objects.requireNonNull(result, "result"); + } + } + + record Failure(Document id, McpError error) implements McpOutcome { + public Failure { + Objects.requireNonNull(error, "error"); + } + } + + enum NoResponse implements McpOutcome { + INSTANCE + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPromptDescriptor.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPromptDescriptor.java new file mode 100644 index 0000000000..6b68488ecf --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPromptDescriptor.java @@ -0,0 +1,8 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +record McpPromptDescriptor(Prompt prompt, McpRemoteClient remoteClient) {} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocol.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocol.java new file mode 100644 index 0000000000..efeaab302b --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocol.java @@ -0,0 +1,136 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Set; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.model.shapes.ShapeType; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Behavior of one MCP protocol version. + * + *

Supported methods and wire features are declared as immutable protocol data. + */ +@SmithyUnstableApi +public sealed interface McpProtocol permits BuiltInProtocol, ExtensionMcpProtocol { + McpProtocolId id(); + + default Set supportedMethods() { + return Set.of(); + } + + default McpProtocolFeatures features() { + return McpProtocolFeatures.NONE; + } + + default ProtocolVersion protocolVersion() { + return ProtocolVersion.parse(id().identifier()); + } + + default void validate(McpCall call, McpRequestContext context) { + var id = call.id(); + if (id != null + && !(id.isType(ShapeType.INTEGER) + || id.isType(ShapeType.LONG) + || id.isType(ShapeType.BIG_INTEGER) + || id.isType(ShapeType.STRING))) { + throw new McpProtocolException(-32602, "Request id is of invalid type " + id.type().name()); + } + if (id == null && !call.method().wireName().startsWith("notifications/")) { + throw new McpProtocolException(-32602, "Requests are expected to have ids"); + } + if (!usesStatelessMetadata()) { + return; + } + + var metadata = call.metadata(); + if (metadata.protocolVersion() == null) { + throw new McpProtocolException( + -32602, + "Missing " + McpWireNames.PROTOCOL_VERSION + " in params._meta"); + } + if (!metadata.protocolVersion().identifier().equals(id().identifier())) { + throw new McpProtocolException( + -32022, + "Unsupported protocol version: " + metadata.protocolVersion().identifier()); + } + var capabilities = metadata.clientCapabilities(); + if (capabilities == null + || !(capabilities.isType(ShapeType.MAP) || capabilities.isType(ShapeType.STRUCTURE))) { + throw new McpProtocolException( + -32602, + "Missing or invalid " + McpWireNames.CLIENT_CAPABILITIES + " in params._meta"); + } + } + + default McpOutcome dispatch(McpCall call, McpOperations operations, McpRequestContext context) { + if (call instanceof McpCall.ExtensionCall extension) { + return executeExtension(extension, context); + } + if (!(call.method() instanceof McpMethod.Standard standard) + || !supportedMethods().contains(standard)) { + throw unsupported(call.method()); + } + + return switch (call) { + case McpCall.Initialize initialize -> operations.initialize(initialize, context); + case McpCall.Ping ping -> operations.ping(ping, context); + case McpCall.Discover discover -> operations.discover(discover, context); + case McpCall.ListTools listTools -> operations.listTools(listTools, context); + case McpCall.CallTool callTool -> operations.callTool(callTool, context); + case McpCall.ListPrompts listPrompts -> operations.listPrompts(listPrompts, context); + case McpCall.GetPrompt getPrompt -> operations.getPrompt(getPrompt, context); + case McpCall.Complete complete -> operations.complete(complete, context); + case McpCall.SetLogLevel setLogLevel -> { + if (setLogLevel.level() == null) { + throw new McpProtocolException(-32602, "Missing or invalid string parameter: level"); + } + yield operations.setLogLevel(setLogLevel, context); + } + case McpCall.ReadResource readResource -> operations.readResource(readResource, context); + case McpCall.Notification ignored -> McpOutcome.NoResponse.INSTANCE; + case McpCall.ExtensionCall ignored -> + throw new IllegalStateException("Extension calls are dispatched before standard calls"); + case McpCall.UnknownCall unknown -> throw unsupported(unknown.method()); + }; + } + + private static

McpOutcome executeExtension( + McpCall.ExtensionCall

extension, + McpRequestContext context + ) { + return extension.extension().execute(extension, context); + } + + default boolean supportsOutputSchema() { + return features().outputSchema(); + } + + default boolean supportsAnnotations() { + return features().annotations(); + } + + default boolean usesStatelessMetadata() { + return features().statelessMetadata(); + } + + default boolean usesHttpMethodHeaders() { + return features().httpMethodHeaders(); + } + + default Document decorateResult( + Document result, + McpMethod method, + McpServerIdentity serverIdentity + ) { + return result; + } + + default McpUnsupportedMethodException unsupported(McpMethod method) { + return new McpUnsupportedMethodException(method, id()); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolException.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolException.java new file mode 100644 index 0000000000..ce4b90df85 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolException.java @@ -0,0 +1,36 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * An MCP protocol-level error with a JSON-RPC error code and optional data. + */ +@SmithyUnstableApi +public final class McpProtocolException extends RuntimeException { + private final int code; + private final Document data; + + public McpProtocolException(int code, String message) { + this(code, message, null); + } + + public McpProtocolException(int code, String message, Document data) { + super(message); + this.code = code; + this.data = data; + } + + public int code() { + return code; + } + + public Document data() { + return data; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolFeatures.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolFeatures.java new file mode 100644 index 0000000000..399fee6983 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolFeatures.java @@ -0,0 +1,23 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Wire and projection features enabled by an MCP protocol. + */ +@SmithyUnstableApi +public record McpProtocolFeatures( + boolean outputSchema, + boolean annotations, + boolean statelessMetadata, + boolean httpMethodHeaders, + boolean statelessResults) { + + public static final McpProtocolFeatures NONE = + new McpProtocolFeatures(false, false, false, false, false); +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolId.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolId.java new file mode 100644 index 0000000000..070e0bdcb7 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolId.java @@ -0,0 +1,26 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Typed registry key for an MCP protocol version. + */ +@SmithyUnstableApi +public record McpProtocolId(String identifier) { + public McpProtocolId { + Objects.requireNonNull(identifier, "identifier"); + if (identifier.isBlank()) { + throw new IllegalArgumentException("MCP protocol identifier must not be blank"); + } + } + + public static McpProtocolId of(String identifier) { + return new McpProtocolId(identifier); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolProvider.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolProvider.java new file mode 100644 index 0000000000..76f46bca64 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolProvider.java @@ -0,0 +1,20 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Collection; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Service-provider interface for discovering external MCP protocols. + * + *

Providers are registered through + * {@code META-INF/services/software.amazon.smithy.java.mcp.server.McpProtocolProvider}. + */ +@SmithyUnstableApi +public interface McpProtocolProvider { + Collection protocols(); +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolRegistry.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolRegistry.java new file mode 100644 index 0000000000..5a68ba7840 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolRegistry.java @@ -0,0 +1,181 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.ServiceLoader; +import software.amazon.smithy.java.core.serde.document.Document; + +/** + * Immutable registry of built-in and extension MCP protocols. + */ +final class McpProtocolRegistry { + private final Map protocols; + private final List supportedIdentifiers; + private final McpProtocol initializationFallbackProtocol; + + private McpProtocolRegistry(Map protocols) { + this.protocols = Collections.unmodifiableMap(new LinkedHashMap<>(protocols)); + supportedIdentifiers = protocols.keySet().stream().map(McpProtocolId::identifier).toList(); + initializationFallbackProtocol = this.protocols.values() + .stream() + .filter(protocol -> protocol.supportedMethods().contains(McpMethod.Standard.INITIALIZE)) + .filter(protocol -> !protocol.usesStatelessMetadata()) + .findFirst() + .orElseThrow(() -> new IllegalStateException( + "No initialization-capable MCP protocol is registered")); + } + + static McpProtocolRegistry create( + Collection additions, + Collection overrides, + boolean discover + ) { + return discover + ? create( + additions, + overrides, + McpProtocolProvider.class.getClassLoader()) + : create(additions, overrides, List.of()); + } + + static McpProtocolRegistry create( + Collection additions, + Collection overrides, + ClassLoader classLoader + ) { + return create( + additions, + overrides, + ServiceLoader.load(McpProtocolProvider.class, classLoader)); + } + + static McpProtocolRegistry create( + Collection additions, + Collection overrides, + Iterable providers + ) { + var candidates = new LinkedHashMap>(); + for (var protocol : BuiltInProtocols.all()) { + addCandidate(candidates, protocol, "built in"); + } + for (var provider : providers) { + Objects.requireNonNull(provider, "MCP protocol provider"); + var provided = Objects.requireNonNull( + provider.protocols(), + () -> "MCP protocol provider returned null: " + provider.getClass().getName()); + for (var protocol : provided) { + addCandidate( + candidates, + protocol, + "SPI provider " + provider.getClass().getName()); + } + } + for (var protocol : additions) { + addCandidate(candidates, protocol, "engine builder"); + } + + var explicitOverrides = new HashMap(); + for (var protocol : overrides) { + Objects.requireNonNull(protocol, "MCP protocol override"); + var previous = explicitOverrides.put(protocol.id(), protocol); + if (previous != null) { + throw new IllegalArgumentException( + "Duplicate MCP protocol override: " + protocol.id().identifier()); + } + } + + var resolved = new LinkedHashMap(); + for (var entry : candidates.entrySet()) { + var override = explicitOverrides.remove(entry.getKey()); + if (override != null) { + resolved.put(entry.getKey(), override); + continue; + } + + var registrations = entry.getValue(); + if (registrations.size() > 1) { + throw conflict(entry.getKey(), registrations); + } + resolved.put(entry.getKey(), registrations.getFirst().protocol()); + } + if (!explicitOverrides.isEmpty()) { + var id = explicitOverrides.keySet().iterator().next(); + throw new IllegalArgumentException( + "Cannot override unregistered MCP protocol: " + id.identifier()); + } + return new McpProtocolRegistry(resolved); + } + + McpProtocol require(ProtocolVersion version) { + var protocol = find(version); + if (protocol != null) { + return protocol; + } + throw new McpProtocolException( + -32022, + "Unsupported protocol version: " + version.identifier(), + Document.of(Map.of( + "requested", + Document.of(version.identifier()), + "supported", + Document.of(supportedIdentifiers.stream() + .map(Document::of) + .toList())))); + } + + McpProtocol find(ProtocolVersion version) { + return protocols.get(McpProtocolId.of(version.identifier())); + } + + McpProtocol defaultProtocol() { + return protocols.get(ProtocolVersion.defaultVersion().id()); + } + + McpProtocol initializationFallbackProtocol() { + return initializationFallbackProtocol; + } + + List supportedIdentifiers() { + return supportedIdentifiers; + } + + private static void addCandidate( + Map> candidates, + McpProtocol protocol, + String source + ) { + Objects.requireNonNull(protocol, "MCP protocol"); + Objects.requireNonNull(protocol.id(), "MCP protocol id"); + candidates.computeIfAbsent(protocol.id(), ignored -> new ArrayList<>()) + .add(new Candidate(protocol, source)); + } + + private static IllegalStateException conflict( + McpProtocolId id, + List candidates + ) { + var sources = candidates.stream() + .map(candidate -> candidate.protocol().getClass().getName() + " from " + candidate.source()) + .sorted() + .toList(); + return new IllegalStateException( + "Conflicting MCP protocol implementations for " + + id.identifier() + + ": " + + String.join(", ", sources) + + ". Use overrideProtocol to select one explicitly."); + } + + private record Candidate(McpProtocol protocol, String source) {} +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteClient.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteClient.java new file mode 100644 index 0000000000..c27b800a78 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteClient.java @@ -0,0 +1,159 @@ +/* + * 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 software.amazon.smithy.java.mcp.model.ListPromptsResult.builder; + +import java.util.List; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import software.amazon.smithy.java.core.schema.SerializableStruct; +import software.amazon.smithy.java.core.schema.ShapeBuilder; +import software.amazon.smithy.java.core.serde.document.Document; +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.ListToolsResult; +import software.amazon.smithy.java.mcp.model.PromptInfo; +import software.amazon.smithy.java.mcp.model.ToolInfo; +import software.amazon.smithy.model.shapes.ShapeType; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Blocking client for a remote MCP server. + * + *

Implementations may use asynchronous I/O internally, but callers observe a single + * blocking exchange operation suitable for execution on virtual threads. + */ +@SmithyUnstableApi +public abstract class McpRemoteClient implements AutoCloseable { + + private static final InternalLogger LOG = InternalLogger.getLogger(McpRemoteClient.class); + private static final AtomicInteger ID_GENERATOR = new AtomicInteger(); + + private final AtomicReference> responseNotificationConsumer = new AtomicReference<>(); + private final AtomicReference> requestNotificationConsumer = new AtomicReference<>(); + private final AtomicReference protocol = + new AtomicReference<>(BuiltInProtocols.protocol(ProtocolVersion.defaultVersion())); + + public List listTools() { + var response = exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.TOOLS_LIST.wireName()) + .id(generateRequestId()) + .jsonrpc("2.0") + .build()); + requireSuccess(response, "listing tools"); + return response.getResult().asShape(ListToolsResult.builder()).getTools().stream().toList(); + } + + public List listPrompts() { + var response = exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.PROMPTS_LIST.wireName()) + .id(generateRequestId()) + .jsonrpc("2.0") + .build()); + requireSuccess(response, "listing prompts"); + return response.getResult().asShape(builder()).getPrompts().stream().toList(); + } + + final void initialize( + Consumer responseNotificationConsumer, + Consumer requestNotificationConsumer, + JsonRpcRequest initializeRequest, + McpProtocol protocol + ) { + var result = Objects.requireNonNull(exchange(initializeRequest), "initialize response"); + requireSuccess(result, "initialization"); + + exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .jsonrpc("2.0") + .build()); + + this.responseNotificationConsumer.set(responseNotificationConsumer); + this.requestNotificationConsumer.set(requestNotificationConsumer); + this.protocol.set(protocol); + } + + protected final ProtocolVersion protocolVersion() { + return protocol.get().protocolVersion(); + } + + protected final McpProtocol protocol() { + return protocol.get(); + } + + /** + * Performs one blocking JSON-RPC exchange. Notifications return {@code null}. + */ + protected abstract JsonRpcResponse exchange(JsonRpcRequest request); + + /** + * Starts resources owned by this client. + */ + public abstract void start(); + + /** + * Stops resources owned by this client. + */ + @Override + public abstract void close(); + + protected final T exchange(String method, ShapeBuilder builder) { + var response = exchange(JsonRpcRequest.builder() + .method(method) + .id(generateRequestId()) + .jsonrpc("2.0") + .build()); + requireSuccess(response, method); + return response.getResult().asShape(builder); + } + + protected final Document generateRequestId() { + return Document.of(ID_GENERATOR.incrementAndGet()); + } + + protected final void notify(JsonRpcResponse response) { + var consumer = responseNotificationConsumer.get(); + if (consumer != null) { + consumer.accept(response); + } + } + + protected final void notify(JsonRpcRequest notification) { + var consumer = requestNotificationConsumer.get(); + if (consumer != null) { + LOG.debug("Forwarding notification to consumer: method={}", notification.getMethod()); + consumer.accept(notification); + } else { + LOG.warn("No request notification consumer set, dropping notification: method={}", + notification.getMethod()); + } + } + + protected static boolean isNotification(Document doc) { + try { + return (doc.isType(ShapeType.STRUCTURE) || doc.isType(ShapeType.MAP)) + && doc.getMember("id") == null + && doc.getMember("method") != null; + } catch (RuntimeException e) { + LOG.warn("Failed to determine whether MCP document is a notification", e); + return false; + } + } + + private static void requireSuccess(JsonRpcResponse response, String action) { + Objects.requireNonNull(response, action + " response"); + if (response.getError() != null) { + throw new McpRemoteException("Remote MCP error during " + action + ": " + + response.getError().getMessage()); + } + } + + public abstract String name(); +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteException.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteException.java new file mode 100644 index 0000000000..f05331b909 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteException.java @@ -0,0 +1,19 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +/** + * Failure while communicating with a remote MCP server. + */ +public final class McpRemoteException extends RuntimeException { + McpRemoteException(String message) { + super(message); + } + + McpRemoteException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestContext.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestContext.java new file mode 100644 index 0000000000..249ced74e1 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestContext.java @@ -0,0 +1,25 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Immutable request-scoped execution context. + */ +@SmithyUnstableApi +public record McpRequestContext( + ProtocolVersion protocolVersion, + McpTransportContext transport, + Context attributes) { + public McpRequestContext { + Objects.requireNonNull(protocolVersion, "protocolVersion"); + transport = transport == null ? McpTransportContext.STDIO : transport; + attributes = attributes == null ? Context.create() : attributes; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestDecoder.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestDecoder.java new file mode 100644 index 0000000000..f3529425b6 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestDecoder.java @@ -0,0 +1,184 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.model.shapes.ShapeType; + +final class McpRequestDecoder { + private static final int INVALID_PARAMS = -32602; + + private final Map> extensions; + + McpRequestDecoder(Map> extensions) { + this.extensions = Map.copyOf(extensions); + } + + McpCall decode(JsonRpcRequest request) { + Objects.requireNonNull(request, "request"); + var params = request.getParams(); + var metadata = decodeMetadata(params); + return switch (McpMethod.parse(request.getMethod())) { + case McpMethod.Standard.INITIALIZE -> decodeInitialize(request, params, metadata); + case McpMethod.Standard.PING -> new McpCall.Ping(request.getId(), metadata); + case McpMethod.Standard.SERVER_DISCOVER -> new McpCall.Discover(request.getId(), metadata); + case McpMethod.Standard.TOOLS_LIST -> + new McpCall.ListTools(request.getId(), optionalString(params, "cursor"), metadata); + case McpMethod.Standard.TOOLS_CALL -> new McpCall.CallTool( + request.getId(), + requiredString(params, "name"), + member(params, "arguments"), + metadata); + case McpMethod.Standard.PROMPTS_LIST -> + new McpCall.ListPrompts(request.getId(), optionalString(params, "cursor"), metadata); + case McpMethod.Standard.PROMPTS_GET -> new McpCall.GetPrompt( + request.getId(), + requiredString(params, "name"), + documentMap(member(params, "arguments")), + metadata); + case McpMethod.Standard.COMPLETION_COMPLETE -> decodeComplete(request, params, metadata); + case McpMethod.Standard.LOGGING_SET_LEVEL -> + new McpCall.SetLogLevel(request.getId(), optionalString(params, "level"), metadata); + case McpMethod.Standard.RESOURCES_READ -> + new McpCall.ReadResource(request.getId(), requiredString(params, "uri"), metadata); + case McpMethod.Standard standard when standard.wireName().startsWith("notifications/") -> + new McpCall.Notification(standard, params, metadata); + case McpMethod.Standard standard -> + new McpCall.UnknownCall( + request.getId(), + new McpMethod.Unknown(standard.wireName()), + params, + metadata); + case McpMethod.Extension extension -> decodeExtension(request, extension, params, metadata); + case McpMethod.Unknown unknown -> { + var extension = extensions.get(unknown.wireName()); + yield extension == null + ? new McpCall.UnknownCall(request.getId(), unknown, params, metadata) + : decodeExtension(request, new McpMethod.Extension(unknown.wireName()), params, metadata); + } + }; + } + + private McpCall.Initialize decodeInitialize( + JsonRpcRequest request, + Document params, + McpMetadata metadata + ) { + var identifier = optionalString(params, "protocolVersion"); + return new McpCall.Initialize( + request.getId(), + ProtocolVersion.parse(identifier), + member(params, "clientInfo"), + member(params, "capabilities"), + metadata); + } + + private McpCall.Complete decodeComplete( + JsonRpcRequest request, + Document params, + McpMetadata metadata + ) { + var reference = member(params, "ref"); + var argument = member(params, "argument"); + return new McpCall.Complete( + request.getId(), + reference == null + ? null + : new McpCall.CompletionReference( + optionalString(reference, "type"), + optionalString(reference, "name")), + argument == null + ? null + : new McpCall.CompletionArgument( + optionalString(argument, "name"), + optionalString(argument, "value")), + metadata); + } + + @SuppressWarnings("unchecked") + private

McpCall.ExtensionCall

decodeExtension( + JsonRpcRequest request, + McpMethod.Extension method, + Document params, + McpMetadata metadata + ) { + var extension = (McpExtensionMethod

) extensions.get(method.wireName()); + if (extension == null) { + throw new IllegalStateException("Unregistered MCP extension: " + method.wireName()); + } + return new McpCall.ExtensionCall<>(request.getId(), extension, extension.decode(params), metadata); + } + + private McpMetadata decodeMetadata(Document params) { + var meta = member(params, "_meta"); + if (meta == null) { + return McpMetadata.EMPTY; + } + if (!isObject(meta)) { + throw invalidParams("params._meta must be an object"); + } + + var values = new HashMap<>(meta.asStringMap()); + var version = removeString(values, McpWireNames.PROTOCOL_VERSION); + var clientInfo = values.remove(McpWireNames.CLIENT_INFO); + var capabilities = values.remove(McpWireNames.CLIENT_CAPABILITIES); + return new McpMetadata( + version == null ? null : ProtocolVersion.parse(version), + clientInfo, + capabilities, + values); + } + + private String removeString(Map values, String name) { + var value = values.remove(name); + if (value == null) { + return null; + } + if (!value.isType(ShapeType.STRING)) { + throw invalidParams(name + " must be a string"); + } + return value.asString(); + } + + private String requiredString(Document document, String name) { + var value = optionalString(document, name); + if (value == null) { + throw invalidParams("Missing or invalid string parameter: " + name); + } + return value; + } + + private String optionalString(Document document, String name) { + var value = member(document, name); + if (value == null) { + return null; + } + if (!value.isType(ShapeType.STRING)) { + throw invalidParams(name + " must be a string"); + } + return value.asString(); + } + + private Document member(Document document, String name) { + return document == null || !isObject(document) ? null : document.getMember(name); + } + + private Map documentMap(Document document) { + return document == null ? Map.of() : Map.copyOf(document.asStringMap()); + } + + private boolean isObject(Document document) { + return document.isType(ShapeType.MAP) || document.isType(ShapeType.STRUCTURE); + } + + private McpProtocolException invalidParams(String message) { + return new McpProtocolException(INVALID_PARAMS, message); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java new file mode 100644 index 0000000000..64ee01def9 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java @@ -0,0 +1,363 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import software.amazon.smithy.ai.McpHeaderTrait; +import software.amazon.smithy.java.core.schema.Schema; +import software.amazon.smithy.java.core.schema.SchemaIndex; +import software.amazon.smithy.java.core.schema.SerializableShape; +import software.amazon.smithy.java.core.schema.TraitKey; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.OneOfTrait; +import software.amazon.smithy.java.mcp.model.JsonArraySchema; +import software.amazon.smithy.java.mcp.model.JsonDocumentSchema; +import software.amazon.smithy.java.mcp.model.JsonObjectSchema; +import software.amazon.smithy.java.mcp.model.JsonOneOfSchema; +import software.amazon.smithy.java.mcp.model.JsonPrimitiveSchema; +import software.amazon.smithy.java.mcp.model.JsonPrimitiveType; +import software.amazon.smithy.java.mcp.model.ToolAnnotations; +import software.amazon.smithy.java.mcp.model.ToolInfo; +import software.amazon.smithy.java.server.Operation; +import software.amazon.smithy.java.server.Service; +import software.amazon.smithy.model.shapes.ShapeId; +import software.amazon.smithy.model.shapes.ShapeType; + +/** + * Converts Smithy operation schemas into canonical MCP tool descriptors. + */ +final class McpSchemaFactory { + private static final TraitKey ONE_OF_TRAIT = TraitKey.get(OneOfTrait.class); + private static final TraitKey MCP_HEADER_TRAIT = TraitKey.get(McpHeaderTrait.class); + private static final List DOCUMENT_TYPES = List.of( + "string", + "number", + "boolean", + "object", + "array", + "null"); + + private final SchemaIndex schemaIndex; + + McpSchemaFactory(SchemaIndex schemaIndex) { + this.schemaIndex = schemaIndex; + } + + Map createTools(Map services) { + var tools = new HashMap(); + for (var entry : services.entrySet()) { + var serverId = entry.getKey(); + var service = entry.getValue(); + for (var operation : service.getAllOperations()) { + var descriptor = createTool(serverId, service, operation); + tools.put(descriptor.info().getName(), descriptor); + } + } + return tools; + } + + McpToolDescriptor createTool(String serverId, Service service, Operation operation) { + var operationSchema = operation.getApiOperation().schema(); + var operationName = operation.name(); + var cache = new HashMap(); + var info = ToolInfo.builder() + .name(operationName) + .description(createDescription(service.schema().id().getName(), operationName, operationSchema)) + .inputSchema(createObjectSchema( + operation.getApiOperation().inputSchema(), + operation.getApiOperation().inputSchema(), + new HashSet<>(), + cache)) + .outputSchema(createObjectSchema( + operation.getApiOperation().outputSchema(), + operation.getApiOperation().outputSchema(), + new HashSet<>(), + cache)) + .annotations(createAnnotations(operationSchema)) + .build(); + return new McpToolDescriptor( + info, + serverId, + new McpToolDescriptor.LocalTarget(operation), + localHeaderParameters(operation)); + } + + private Map localHeaderParameters(Operation operation) { + var result = new HashMap(); + for (var member : operation.getApiOperation().inputSchema().members()) { + var trait = member.getTrait(MCP_HEADER_TRAIT); + if (trait != null && trait.getValue().matches("[A-Za-z0-9][A-Za-z0-9_-]*")) { + result.put(member.memberName(), trait.getValue()); + } + } + return Map.copyOf(result); + } + + private ToolAnnotations createAnnotations(Schema operationSchema) { + boolean readOnly = operationSchema.hasTrait(TraitKey.READ_ONLY_TRAIT); + boolean idempotent = operationSchema.hasTrait(TraitKey.IDEMPOTENT_TRAIT); + if (!readOnly && !idempotent) { + return null; + } + var builder = ToolAnnotations.builder(); + if (readOnly) { + builder.readOnlyHint(true); + } + if (idempotent) { + builder.idempotentHint(true); + } + return builder.build(); + } + + private JsonObjectSchema createObjectSchema( + Schema member, + Schema target, + Set visited, + Map cache + ) { + var targetId = target.id(); + var cached = cache.get(targetId); + if (cached != null) { + return (JsonObjectSchema) withDescription(cached, memberDescription(member)); + } + if (!visited.add(targetId)) { + return JsonObjectSchema.builder().build(); + } + + var properties = new HashMap(); + var required = new ArrayList(); + for (var child : target.members()) { + if (child.hasTrait(TraitKey.REQUIRED_TRAIT)) { + required.add(child.memberName()); + } + properties.put(child.memberName(), Document.of(createMemberSchema(child, visited, cache))); + } + visited.remove(targetId); + + var result = JsonObjectSchema.builder() + .properties(properties) + .required(required) + .build(); + cache.put(targetId, result); + return (JsonObjectSchema) withDescription(result, memberDescription(member)); + } + + private JsonArraySchema createArraySchema( + Schema member, + Schema target, + Set visited, + Map cache + ) { + var items = createMemberSchema(target.listMember(), visited, cache); + var itemDocument = target.hasTrait(TraitKey.SPARSE_TRAIT) + ? Document.of(Map.of( + "anyOf", + Document.of(List.of( + Document.of(items), + Document.of(Map.of("type", Document.of("null"))))))) + : Document.of(items); + return JsonArraySchema.builder() + .description(memberDescription(member)) + .items(itemDocument) + .build(); + } + + private JsonPrimitiveSchema createPrimitiveSchema(Schema member) { + var type = switch (member.type()) { + case BYTE, SHORT, INTEGER, INT_ENUM, LONG, FLOAT, DOUBLE -> JsonPrimitiveType.NUMBER; + case ENUM, BLOB, STRING, BIG_DECIMAL, BIG_INTEGER, TIMESTAMP -> JsonPrimitiveType.STRING; + case BOOLEAN -> JsonPrimitiveType.BOOLEAN; + default -> throw new IllegalArgumentException(member + " is not a primitive type"); + }; + + var builder = JsonPrimitiveSchema.builder() + .type(type) + .description(memberDescription(member)); + var header = member.getTrait(MCP_HEADER_TRAIT); + if (header != null) { + builder.mcpHeader(header.getValue()); + } + if (member.type() == ShapeType.TIMESTAMP) { + builder.format("date-time"); + } + + List enumValues = switch (member.type()) { + case ENUM, STRING -> member.stringEnumValues().stream().map(Document::of).toList(); + case INT_ENUM -> member.intEnumValues().stream().map(Document::of).toList(); + default -> List.of(); + }; + if (!enumValues.isEmpty()) { + builder.enumValues(enumValues); + } + return builder.build(); + } + + private SerializableShape createDocumentSchema( + Schema member, + Set visited, + Map cache + ) { + var target = member.isMember() ? member.memberTarget() : member; + var oneOf = target.getTrait(ONE_OF_TRAIT); + if (oneOf == null) { + return JsonDocumentSchema.builder() + .type(DOCUMENT_TYPES) + .description(memberDescription(member)) + .build(); + } + return createOneOfSchema(oneOf, member, visited, cache); + } + + private SerializableShape createOneOfSchema( + OneOfTrait oneOf, + Schema documentMember, + Set visited, + Map cache + ) { + var targetId = (documentMember.isMember() ? documentMember.memberTarget() : documentMember).id(); + var cached = cache.get(targetId); + if (cached != null) { + return withDescription(cached, memberDescription(documentMember)); + } + if (!visited.add(targetId)) { + return JsonObjectSchema.builder().build(); + } + + var variants = new ArrayList(); + for (var definition : oneOf.getMembers()) { + var target = schemaIndex.getSchema(definition.getTarget()); + variants.add(createUnionVariant( + definition.getName(), + createObjectSchema(target, target, visited, cache))); + } + visited.remove(targetId); + + var result = JsonOneOfSchema.builder().oneOf(variants).build(); + cache.put(targetId, result); + return withDescription(result, memberDescription(documentMember)); + } + + private SerializableShape createUnionSchema( + Schema member, + Schema target, + Set visited, + Map cache + ) { + var targetId = target.id(); + var cached = cache.get(targetId); + if (cached != null) { + return withDescription(cached, memberDescription(member)); + } + if (!visited.add(targetId)) { + return JsonObjectSchema.builder().build(); + } + + var variants = new ArrayList(); + for (var child : target.members()) { + variants.add(createUnionVariant( + child.memberName(), + createMemberSchema(child, visited, cache))); + } + visited.remove(targetId); + + var result = JsonOneOfSchema.builder().oneOf(variants).build(); + cache.put(targetId, result); + return withDescription(result, memberDescription(member)); + } + + private SerializableShape createMemberSchema( + Schema member, + Set visited, + Map cache + ) { + return switch (member.type()) { + case LIST, SET -> createArraySchema(member, member.memberTarget(), visited, cache); + case MAP -> createMapSchema(member, member.memberTarget(), visited, cache); + case STRUCTURE -> createObjectSchema(member, member.memberTarget(), visited, cache); + case UNION -> createUnionSchema(member, member.memberTarget(), visited, cache); + case DOCUMENT -> createDocumentSchema(member, visited, cache); + default -> createPrimitiveSchema(member); + }; + } + + private JsonObjectSchema createMapSchema( + Schema member, + Schema target, + Set visited, + Map cache + ) { + var value = createMemberSchema(target.mapValueMember(), visited, cache); + var additionalProperties = target.hasTrait(TraitKey.SPARSE_TRAIT) + ? Document.of(Map.of( + "anyOf", + Document.of(List.of( + Document.of(value), + Document.of(Map.of("type", Document.of("null"))))))) + : Document.of(value); + return JsonObjectSchema.builder() + .description(memberDescription(member)) + .additionalProperties(additionalProperties) + .build(); + } + + private static Document createUnionVariant(String memberName, SerializableShape memberSchema) { + return Document.of(JsonObjectSchema.builder() + .properties(Map.of(memberName, Document.of(memberSchema))) + .required(List.of(memberName)) + .additionalProperties(Document.of(false)) + .build()); + } + + private static String memberDescription(Schema schema) { + String description = null; + var trait = schema.isMember() + ? schema.getDirectTrait(TraitKey.DOCUMENTATION_TRAIT) + : schema.getTrait(TraitKey.DOCUMENTATION_TRAIT); + if (trait != null) { + description = trait.getValue(); + } + if (schema.isMember()) { + var targetDescription = memberDescription(schema.memberTarget()); + if (description != null && targetDescription != null) { + description = appendSentences(description, targetDescription); + } else if (targetDescription != null) { + description = targetDescription; + } + } + return description; + } + + private static String createDescription(String serviceName, String operationName, Schema schema) { + var documentation = schema.getTrait(TraitKey.DOCUMENTATION_TRAIT); + return documentation == null + ? "This tool invokes %s API of %s.".formatted(operationName, serviceName) + : documentation.getValue(); + } + + private static String appendSentences(String first, String second) { + first = first.trim(); + if (!first.endsWith(".")) { + first += ". "; + } + return first + second; + } + + private static SerializableShape withDescription(SerializableShape schema, String description) { + if (description == null) { + return schema; + } + return switch (schema) { + case JsonObjectSchema object -> object.toBuilder().description(description).build(); + case JsonOneOfSchema oneOf -> oneOf.toBuilder().description(description).build(); + default -> schema; + }; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServer.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServer.java deleted file mode 100644 index c98d024f6a..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServer.java +++ /dev/null @@ -1,167 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; -import java.nio.charset.StandardCharsets; -import java.util.Scanner; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CountDownLatch; -import software.amazon.smithy.java.core.schema.SerializableStruct; -import software.amazon.smithy.java.io.ByteBufferUtils; -import software.amazon.smithy.java.json.JsonCodec; -import software.amazon.smithy.java.json.JsonSettings; -import software.amazon.smithy.java.logging.InternalLogger; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; -import software.amazon.smithy.java.server.Server; -import software.amazon.smithy.java.server.Service; -import software.amazon.smithy.utils.SmithyUnstableApi; - -@SmithyUnstableApi -public final class McpServer implements Server { - - private static final InternalLogger LOG = InternalLogger.getLogger(McpServer.class); - - private static final JsonCodec CODEC = JsonCodec.builder() - .settings(JsonSettings.builder() - .serializeTypeInDocuments(false) - .useJsonName(true) - .build()) - .build(); - - private final McpService mcpService; - private final Thread listener; - private final InputStream is; - private final OutputStream os; - private final CountDownLatch done = new CountDownLatch(1); - private volatile ProtocolVersion protocolVersion; - - McpServer(McpServerBuilder builder) { - this.mcpService = builder.mcpService; - this.is = builder.is; - this.os = builder.os; - this.listener = new Thread(() -> { - try { - this.listen(); - } catch (Exception e) { - LOG.error("Error handling request", e); - } finally { - done.countDown(); - } - }); - listener.setName("stdio-dispatcher"); - listener.setDaemon(true); - } - - private void listen() { - var scan = new Scanner(is, StandardCharsets.UTF_8); - while (scan.hasNextLine()) { - var line = scan.nextLine(); - try { - var jsonRequest = CODEC.deserializeShape(line, JsonRpcRequest.builder()); - handleRequest(jsonRequest); - } catch (Exception e) { - LOG.error("Error decoding request", e); - } - } - } - - private void handleRequest(JsonRpcRequest req) { - // For StdIO transport, protocol version is only sent in initialize request - // Extract and store it for future requests - if ("initialize".equals(req.getMethod())) { - var maybeVersion = req.getParams().getMember("protocolVersion"); - if (maybeVersion == null) { - this.protocolVersion = ProtocolVersion.defaultVersion(); - } else { - this.protocolVersion = ProtocolVersion.version(maybeVersion.asString()); - } - } - - var response = mcpService.handleRequest(req, this::writeStructToOutput, protocolVersion); - if (response != null) { - writeStructToOutput(response); - } - } - - private static final byte[] TOOLS_CHANGED = """ - {"jsonrpc":"2.0","method":"notifications/tools/list_changed"} - """.getBytes(StandardCharsets.UTF_8); // newline is important here - - public void refreshTools() { - try { - synchronized (os) { - os.write(TOOLS_CHANGED); - os.flush(); - } - } catch (IOException e) { - LOG.error("Failed to flush tools changed notification", e); - } - } - - public void addNewService(String id, Service service) { - mcpService.addNewService(id, service); - refreshTools(); - } - - public void addNewProxy(McpServerProxy mcpServerProxy) { - mcpService.addNewProxy(mcpServerProxy, this::writeStructToOutput); - refreshTools(); - } - - public boolean containsMcpServer(String id) { - return mcpService.containsMcpServer(id); - } - - private void writeStructToOutput(SerializableStruct shape) { - synchronized (os) { - var bytes = CODEC.serialize(shape); - try { - if (bytes.hasArray()) { - os.write(bytes.array(), bytes.arrayOffset() + bytes.position(), bytes.remaining()); - } else { - os.write(ByteBufferUtils.getBytes(bytes)); - } - os.write('\n'); - os.flush(); - } catch (Exception e) { - LOG.error("Error writing to output", e); - } - } - } - - @Override - public void start() { - // Set up notification writer for proxies - mcpService.setNotificationWriter(this::writeStructToOutput); - - // Initialize proxies - mcpService.startProxies(); - - // Start the listener thread - listener.start(); - } - - @Override - public CompletableFuture shutdown() { - return CompletableFuture.allOf( - mcpService.getProxies() - .values() - .stream() - .map(McpServerProxy::shutdown) - .toArray(CompletableFuture[]::new)); - } - - public void awaitCompletion() throws InterruptedException { - done.await(); - } - - public static McpServerBuilder builder() { - return new McpServerBuilder(); - } -} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerBuilder.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerBuilder.java deleted file mode 100644 index 5777859c0b..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerBuilder.java +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import java.io.InputStream; -import java.io.OutputStream; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import software.amazon.smithy.java.server.Server; -import software.amazon.smithy.java.server.Service; -import software.amazon.smithy.utils.SmithyUnstableApi; - -@SmithyUnstableApi -public final class McpServerBuilder { - - InputStream is; - OutputStream os; - Map services = new HashMap<>(); - List proxyList = new ArrayList<>(); - McpServerInterceptor interceptor; - String name; - String version; - ToolFilter toolFilter = (server, tool) -> true; - McpMetricsObserver metricsObserver; - McpService mcpService; - - McpServerBuilder() {} - - public McpServerBuilder stdio() { - this.is = System.in; - this.os = System.out; - return this; - } - - public McpServerBuilder input(InputStream is) { - this.is = is; - return this; - } - - public McpServerBuilder output(OutputStream os) { - this.os = os; - return this; - } - - public McpServerBuilder name(String name) { - this.name = name; - return this; - } - - public McpServerBuilder version(String version) { - this.version = version; - return this; - } - - public Server build() { - validate(); - // Create McpService before building McpServer - var builder = McpService.builder() - .services(services) - .proxyList(proxyList) - .name(name != null ? name : "mcp-server") - .toolFilter(toolFilter) - .metricsObserver(metricsObserver); - - if (version != null) { - builder.version(version); - } - - if (interceptor != null) { - builder.interceptor(interceptor); - } - - this.mcpService = builder.build(); - return new McpServer(this); - } - - public McpServerBuilder addService(String id, Service service) { - services.put(id, service); - return this; - } - - public McpServerBuilder addService(Map services) { - this.services.putAll(services); - return this; - } - - public McpServerBuilder addService(McpServerProxy... proxy) { - proxyList.addAll(Arrays.asList(proxy)); - return this; - } - - public McpServerBuilder toolFilter(ToolFilter filter) { - this.toolFilter = filter; - return this; - } - - public McpServerBuilder metricsObserver(McpMetricsObserver observer) { - this.metricsObserver = observer; - return this; - } - - /** - * Sets the server interceptor. Use {@link McpServerInterceptor#chain(List)} to compose - * multiple interceptors into one. - * - * @see McpServerInterceptor for hook descriptions and the execution lifecycle - */ - public McpServerBuilder interceptor(McpServerInterceptor interceptor) { - this.interceptor = Objects.requireNonNull(interceptor, "interceptor"); - return this; - } - - private void validate() { - Objects.requireNonNull(is, "MCP server input stream is required"); - Objects.requireNonNull(os, "MCP server output stream is required"); - if (services.isEmpty() && proxyList.isEmpty()) { - throw new IllegalArgumentException("MCP server requires at least one service or proxy"); - } - } -} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerIdentity.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerIdentity.java new file mode 100644 index 0000000000..2c04dba1fd --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerIdentity.java @@ -0,0 +1,20 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Identity advertised by an MCP server. + */ +@SmithyUnstableApi +public record McpServerIdentity(String name, String version) { + public McpServerIdentity { + Objects.requireNonNull(name, "name"); + Objects.requireNonNull(version, "version"); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptor.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptor.java deleted file mode 100644 index e2c056441b..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptor.java +++ /dev/null @@ -1,217 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import java.util.List; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; -import software.amazon.smithy.java.mcp.model.JsonRpcResponse; -import software.amazon.smithy.utils.SmithyUnstableApi; - -/** - * Interceptor for MCP server request processing. Interceptors inject code into the - * {@link McpService} request execution pipeline via hooks at specific stages. - * - *

Hooks are either "read" hooks (observe in-flight data) or "modify" hooks (transform - * in-flight data). All hooks have default no-op implementations; override only the hooks - * you need. - * - *

Execution lifecycle

- * - *

For every request: - *

    - *
  1. {@link #readBeforeExecution} — observe the incoming request
    2. - *
  2. {@link #modifyBeforeExecution} — optionally transform the request
    3. - *
  3. For {@code tools/call} requests only: - *
      - *
    1. {@link #readBeforeToolCall} — observe before tool dispatch
      2. - *
    2. {@link #modifyBeforeToolCall} — optionally transform the request
      3. - *
    3. Tool dispatch (local or proxy)
      4. - *
    4. {@link #readAfterToolCall} — observe the tool result
      5. - *
    5. {@link #modifyAfterToolCall} — optionally transform the response
      6. - *
    - *
    4. - *
  4. {@link #readAfterExecution} — observe the final result (ALWAYS fires)
    5. - *
  5. {@link #modifyAfterExecution} — optionally transform the final response
    6. - *
- * - *

Error handling

- * - *

Any hook may throw a {@link RuntimeException}. When a hook throws, remaining hooks - * in that stage are skipped, and execution jumps to the after-execution hooks with the - * error. The {@code readAfterExecution} and {@code modifyAfterExecution} hooks ALWAYS fire, - * ensuring cleanup and telemetry logic runs regardless of errors. - * - *

Async tool calls

- * - *

For proxy tool calls, the after-tool-call and after-execution hooks fire on the - * thread that receives the proxy response, not the original request thread. Hook - * implementations must be thread-safe. - * - *

Example: telemetry

- * 
{@code
- * public class TelemetryInterceptor implements McpServerInterceptor {
- *     private static final Context.Key START = Context.key("start");
- *
- *     @Override
- *     public void readBeforeExecution(McpExecutionHook hook) {
- *         hook.context().put(START, System.nanoTime());
- *     }
- *
- *     @Override
- *     public void readAfterExecution(McpExecutionHook hook,
- *             JsonRpcResponse response, RuntimeException error) {
- *         long duration = System.nanoTime() - hook.context().get(START);
- *         emitMetrics(hook.request().getMethod(), duration, error == null);
- *     }
- * }
- * }
- * - *

Example: access control

- * 
{@code
- * public class AccessControlInterceptor implements McpServerInterceptor {
- *     @Override
- *     public void readBeforeToolCall(McpToolCallHook hook) {
- *         if (isBlocked(hook.toolName(), hook.serverId())) {
- *             throw new RuntimeException("Access denied: " + hook.toolName());
- *         }
- *     }
- * }
- * }
- */ -@SmithyUnstableApi -public interface McpServerInterceptor { - - /** - * An interceptor that does nothing. - */ - McpServerInterceptor NOOP = new McpServerInterceptor() {}; - - /** - * Combines multiple interceptors into a single interceptor that invokes each one - * in order. Hooks are called sequentially on each interceptor in list order. - * - * @param interceptors The interceptors to compose. - * @return A single interceptor that delegates to all provided interceptors. - */ - static McpServerInterceptor chain(List interceptors) { - return switch (interceptors.size()) { - case 0 -> NOOP; - case 1 -> interceptors.get(0); - default -> new McpServerInterceptorChain(List.copyOf(interceptors)); - }; - } - - /** - * Combines multiple interceptors into a single interceptor that invokes each one - * in order. Convenience overload of {@link #chain(List)}. - * - * @param interceptors The interceptors to compose. - * @return A single interceptor that delegates to all provided interceptors. - */ - static McpServerInterceptor chain(McpServerInterceptor... interceptors) { - return chain(List.of(interceptors)); - } - - // --- Execution-level hooks (fire for all requests) --- - - /** - * Called when a request is received, before any dispatch logic. - * - * @param hook Execution hook data containing the request, protocol version, and context. - */ - default void readBeforeExecution(McpExecutionHook hook) {} - - /** - * Called before dispatch. Can return a modified request. - * - * @param hook Execution hook data. - * @return The request to dispatch, or {@code hook.request()} to pass through unmodified. - */ - default JsonRpcRequest modifyBeforeExecution(McpExecutionHook hook) { - return hook.request(); - } - - /** - * Called when execution completes. ALWAYS fires, even if an earlier hook threw. - * - * @param hook Execution hook data. - * @param response The response, or {@code null} for notifications and async proxy calls - * still in flight. - * @param error The error if one occurred, or {@code null} on success. - */ - default void readAfterExecution(McpExecutionHook hook, JsonRpcResponse response, RuntimeException error) {} - - /** - * Called when execution completes. Can modify the response or handle errors. - * ALWAYS fires, even if an earlier hook threw. - * - * @param hook Execution hook data. - * @param response The response, or {@code null} for notifications. - * @param error The error if one occurred, or {@code null} on success. - * @return The final response. - * @throws RuntimeException to propagate or replace the error. - */ - default JsonRpcResponse modifyAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, - RuntimeException error - ) { - if (error != null) { - throw error; - } - return response; - } - - // --- Tool-level hooks (fire only for tools/call) --- - - /** - * Called before a tool is invoked. - * - * @param hook Tool call hook data containing tool name, server ID, and proxy status. - */ - default void readBeforeToolCall(McpToolCallHook hook) {} - - /** - * Called before a tool is invoked. Can return a modified request. - * - * @param hook Tool call hook data. - * @return The request to use for tool invocation, or {@code hook.request()} to pass - * through unmodified. - */ - default JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { - return hook.request(); - } - - /** - * Called after a tool completes. For proxy tools, this fires on the callback thread. - * - * @param hook Tool call hook data. - * @param response The tool call response. - * @param error The error if one occurred, or {@code null} on success. - */ - default void readAfterToolCall(McpToolCallHook hook, JsonRpcResponse response, RuntimeException error) {} - - /** - * Called after a tool completes. Can modify the response or handle errors. - * For proxy tools, this fires on the callback thread. - * - * @param hook Tool call hook data. - * @param response The tool call response. - * @param error The error if one occurred, or {@code null} on success. - * @return The response to return. - * @throws RuntimeException to propagate or replace the error. - */ - default JsonRpcResponse modifyAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, - RuntimeException error - ) { - if (error != null) { - throw error; - } - return response; - } -} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptorChain.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptorChain.java deleted file mode 100644 index 0af3203923..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptorChain.java +++ /dev/null @@ -1,136 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import java.util.List; -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.utils.SmithyUnstableApi; - -/** - * Composes multiple {@link McpServerInterceptor} instances into a single interceptor - * that delegates to each one in order. - */ -@SmithyUnstableApi -final class McpServerInterceptorChain implements McpServerInterceptor { - - private static final InternalLogger LOGGER = InternalLogger.getLogger(McpServerInterceptorChain.class); - private final McpServerInterceptor[] interceptors; - - McpServerInterceptorChain(List interceptors) { - this.interceptors = interceptors.toArray(McpServerInterceptor[]::new); - } - - @Override - public void readBeforeExecution(McpExecutionHook hook) { - RuntimeException error = null; - for (var interceptor : interceptors) { - try { - interceptor.readBeforeExecution(hook); - } catch (RuntimeException e) { - error = swapError("readBeforeExecution", error, e); - } - } - if (error != null) { - throw error; - } - } - - @Override - public JsonRpcRequest modifyBeforeExecution(McpExecutionHook hook) { - var current = hook; - for (var interceptor : interceptors) { - var req = interceptor.modifyBeforeExecution(current); - current = current.withRequest(req); - } - return current.request(); - } - - @Override - public void readAfterExecution(McpExecutionHook hook, JsonRpcResponse response, RuntimeException error) { - for (var interceptor : interceptors) { - try { - interceptor.readAfterExecution(hook, response, error); - } catch (RuntimeException e) { - error = swapError("readAfterExecution", error, e); - } - } - // Always throw the error even if it's the original error. - if (error != null) { - throw error; - } - } - - @Override - public JsonRpcResponse modifyAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, - RuntimeException error - ) { - for (var interceptor : interceptors) { - response = interceptor.modifyAfterExecution(hook, response, error); - error = null; - } - return response; - } - - @Override - public void readBeforeToolCall(McpToolCallHook hook) { - RuntimeException error = null; - for (var interceptor : interceptors) { - try { - interceptor.readBeforeToolCall(hook); - } catch (RuntimeException e) { - error = swapError("readBeforeToolCall", error, e); - } - } - if (error != null) { - throw error; - } - } - - @Override - public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { - var current = hook; - for (var interceptor : interceptors) { - var req = interceptor.modifyBeforeToolCall(current); - current = current.withRequest(req); - } - return current.request(); - } - - @Override - public void readAfterToolCall(McpToolCallHook hook, JsonRpcResponse response, RuntimeException error) { - for (var interceptor : interceptors) { - try { - interceptor.readAfterToolCall(hook, response, error); - } catch (RuntimeException e) { - error = swapError("readAfterToolCall", error, e); - } - } - // Always throw the error even if it's the original error. - if (error != null) { - throw error; - } - } - - @Override - public JsonRpcResponse modifyAfterToolCall(McpToolCallHook hook, JsonRpcResponse response, RuntimeException error) { - for (var interceptor : interceptors) { - response = interceptor.modifyAfterToolCall(hook, response, error); - error = null; - } - return response; - } - - private static RuntimeException swapError(String hook, RuntimeException oldE, RuntimeException newE) { - if (oldE != null && oldE != newE) { - LOGGER.trace("Replacing error after {}: {}", hook, newE.getClass().getName(), newE.getMessage()); - } - return newE; - } -} 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 deleted file mode 100644 index 22e61867f2..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java +++ /dev/null @@ -1,174 +0,0 @@ -/* - * 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 software.amazon.smithy.java.mcp.model.ListPromptsResult.builder; - -import java.util.List; -import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -import software.amazon.smithy.java.core.schema.SerializableStruct; -import software.amazon.smithy.java.core.schema.ShapeBuilder; -import software.amazon.smithy.java.core.serde.document.Document; -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.ListToolsResult; -import software.amazon.smithy.java.mcp.model.PromptInfo; -import software.amazon.smithy.java.mcp.model.ToolInfo; -import software.amazon.smithy.model.shapes.ShapeType; -import software.amazon.smithy.utils.SmithyUnstableApi; - -@SmithyUnstableApi -public abstract class McpServerProxy { - - private static final InternalLogger LOG = InternalLogger.getLogger(McpServerProxy.class); - private static final AtomicInteger ID_GENERATOR = new AtomicInteger(0); - - 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(); - } - - public List listPrompts() { - JsonRpcRequest request = JsonRpcRequest.builder() - .method("prompts/list") - .id(generateRequestId()) - .jsonrpc("2.0") - .build(); - return rpc(request).thenApply(response -> { - if (response.getError() != null) { - throw new RuntimeException("Error listing prompts: " + response.getError().getMessage()); - } - return response.getResult() - .asShape(builder()) - .getPrompts() - .stream() - .toList(); - }).join(); - } - - public void initialize( - Consumer notificationConsumer, - Consumer requestNotificationConsumer, - JsonRpcRequest initializeRequest, - ProtocolVersion protocolVersion - ) { - - var result = Objects.requireNonNull(rpc(initializeRequest).join()); - if (result.getError() != null) { - throw new RuntimeException("Error during initialization: " + result.getError().getMessage()); - } - - // Send the initialized notification per MCP protocol spec - JsonRpcRequest initializedNotification = JsonRpcRequest.builder() - .method("notifications/initialized") - .jsonrpc("2.0") - .build(); - rpc(initializedNotification); - - this.notificationConsumer.set(notificationConsumer); - this.requestNotificationConsumer.set(requestNotificationConsumer); - this.protocolVersion.set(protocolVersion); - } - - protected final ProtocolVersion getProtocolVersion() { - return protocolVersion.get(); - } - - protected abstract CompletableFuture rpc(JsonRpcRequest request); - - protected abstract void start(); - - protected abstract CompletableFuture shutdown(); - - protected CompletableFuture rpc(String method, ShapeBuilder builder) { - JsonRpcRequest request = JsonRpcRequest.builder() - .method(method) - .id(generateRequestId()) - .jsonrpc("2.0") - .build(); - - return rpc(request).thenApply(response -> { - if (response.getError() != null) { - throw new RuntimeException("Error in RPC call: " + response.getError().getMessage()); - } - return response.getResult().asShape(builder); - }); - } - - // Generate a unique request ID for each RPC call - protected Document generateRequestId() { - return Document.of(ID_GENERATOR.incrementAndGet()); - } - - protected void notify(JsonRpcResponse response) { - var nc = notificationConsumer.get(); - if (nc != null) { - nc.accept(response); - } - } - - /** - * Forwards a notification request by converting it to a response format. - * Notifications have a method field but no id. - */ - protected void notify(JsonRpcRequest notification) { - var rnc = requestNotificationConsumer.get(); - if (rnc != null) { - LOG.debug("Forwarding notification to consumer: method={}", notification.getMethod()); - rnc.accept(notification); - } else { - LOG.warn("No request notification consumer set, dropping notification: method={}", - notification.getMethod()); - } - } - - /** - * Determines if a Document represents a notification (has "method" but no "id") - * rather than a response (has "id"). - * - * - Responses have an "id" field at the top level - * - Notifications have a "method" field but no "id" field at the top level - */ - protected static boolean isNotification(Document doc) { - try { - if (!doc.isType(ShapeType.STRUCTURE) && !doc.isType(ShapeType.MAP)) { - return false; - } - - // If it has a "method" field but no "id", it's a notification - return doc.getMember("id") == null && doc.getMember("method") != null; - } catch (Exception e) { - LOG.warn("Failed to determine if notification from Document", e); - return false; - } - } - - public abstract String name(); -} 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 deleted file mode 100644 index 04e15fb914..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java +++ /dev/null @@ -1,1473 +0,0 @@ -/* - * 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 software.amazon.smithy.java.core.serde.TimestampFormatter.Prelude.DATE_TIME; -import static software.amazon.smithy.java.core.serde.TimestampFormatter.Prelude.EPOCH_SECONDS; -import static software.amazon.smithy.java.core.serde.TimestampFormatter.Prelude.HTTP_DATE; -import static software.amazon.smithy.java.mcp.server.PromptLoader.normalize; - -import java.io.PrintWriter; -import java.io.StringWriter; -import java.math.BigDecimal; -import java.math.BigInteger; -import java.util.ArrayList; -import java.util.Base64; -import java.util.HashMap; -import java.util.HashSet; -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.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; -import software.amazon.smithy.java.core.schema.SerializableShape; -import software.amazon.smithy.java.core.schema.TraitKey; -import software.amazon.smithy.java.core.serde.document.Document; -import software.amazon.smithy.java.framework.model.ValidationException; -import software.amazon.smithy.java.io.ByteBufferUtils; -import software.amazon.smithy.java.json.JsonCodec; -import software.amazon.smithy.java.json.JsonSettings; -import software.amazon.smithy.java.logging.InternalLogger; -import software.amazon.smithy.java.mcp.OneOfTrait; -import software.amazon.smithy.java.mcp.model.CallToolResult; -import software.amazon.smithy.java.mcp.model.Capabilities; -import software.amazon.smithy.java.mcp.model.InitializeResult; -import software.amazon.smithy.java.mcp.model.JsonArraySchema; -import software.amazon.smithy.java.mcp.model.JsonDocumentSchema; -import software.amazon.smithy.java.mcp.model.JsonObjectSchema; -import software.amazon.smithy.java.mcp.model.JsonOneOfSchema; -import software.amazon.smithy.java.mcp.model.JsonPrimitiveSchema; -import software.amazon.smithy.java.mcp.model.JsonPrimitiveType; -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.Prompts; -import software.amazon.smithy.java.mcp.model.ServerInfo; -import software.amazon.smithy.java.mcp.model.TextContent; -import software.amazon.smithy.java.mcp.model.ToolAnnotations; -import software.amazon.smithy.java.mcp.model.ToolInfo; -import software.amazon.smithy.java.mcp.model.Tools; -import software.amazon.smithy.java.server.Operation; -import software.amazon.smithy.java.server.Service; -import software.amazon.smithy.model.shapes.ShapeId; -import software.amazon.smithy.model.shapes.ShapeType; -import software.amazon.smithy.utils.SmithyUnstableApi; - -/** - * Core MCP service that handles JSON-RPC requests and returns responses. - * This class is responsible for processing MCP protocol logic independently - * of transport concerns. - */ -@SmithyUnstableApi -public final class McpService { - - private static final InternalLogger LOG = InternalLogger.getLogger(McpService.class); - private static final Context.Key ASYNC_DISPATCH = Context.key("mcp.asyncDispatch"); - private static final int METHOD_NOT_FOUND_ERROR_CODE = -32601; - - private static final JsonCodec CODEC = JsonCodec.builder() - .settings(JsonSettings.builder() - .serializeTypeInDocuments(false) - .useJsonName(true) - .build()) - .build(); - - private static final TraitKey ONE_OF_TRAIT = TraitKey.get(OneOfTrait.class); - - private final Map tools; - private final Map prompts; - 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; - - McpService( - Map services, - List proxyList, - String name, - String version, - ToolFilter toolFilter, - McpMetricsObserver metricsObserver, - McpServerInterceptor interceptor - ) { - this.services = 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.serviceName = name; - this.version = version; - this.proxies = proxyList.stream().collect(Collectors.toMap(McpServerProxy::name, p -> p)); - this.toolFilter = toolFilter; - this.metricsObserver = metricsObserver; - this.interceptor = interceptor; - } - - /** - * Handles a JSON-RPC request, invoking interceptor hooks at each stage of the pipeline. - * - *

Responses are delivered through one of two channels: - *

    - *
  • Synchronous (return value): For most requests, the response is returned directly.
    • - *
  • Asynchronous (callback): For proxy tool calls, returns {@code null} and the callback - * is invoked when the proxy responds.
    • - *
  • Neither: For notifications, returns {@code null} and the callback is never - * invoked. Requests with unknown methods receive a -32601 (Method not found) error.
    • - *
- * - * @param req The JSON-RPC request to handle - * @param asyncResponseCallback Callback for async responses (used for proxy calls) - * @param protocolVersion The protocol version for this request (may be null) - * @return The response for synchronous operations, or null for async/notification operations - */ - public JsonRpcResponse handleRequest( - JsonRpcRequest req, - Consumer asyncResponseCallback, - ProtocolVersion protocolVersion - ) { - // Zero-interceptor fast path: skip Context creation, hook allocation, and all hook invocations. - if (interceptor == McpServerInterceptor.NOOP) { - return handleRequestDirect(req, asyncResponseCallback, protocolVersion); - } - - var hook = new McpExecutionHook(req, protocolVersion, Context.create()); - JsonRpcResponse response = null; - RuntimeException caughtError = null; - - try { - var currentReq = fireBeforeExecution(hook); - hook = hook.withRequest(currentReq); - - // Dispatch - validate(currentReq); - var method = currentReq.getMethod(); - response = switch (method) { - case "initialize" -> handleInitialize(currentReq); - case "ping" -> handlePing(currentReq); - default -> { - initializeProxies(rpcResponse -> {}); - yield switch (method) { - case "prompts/list" -> handlePromptsList(currentReq); - case "prompts/get" -> handlePromptsGet(currentReq); - case "tools/list" -> handleToolsList(currentReq, protocolVersion); - case "tools/call" -> - handleToolsCall(currentReq, asyncResponseCallback, protocolVersion, hook); - default -> methodNotFound(currentReq); - }; - } - }; - if (Boolean.TRUE.equals(hook.context().get(ASYNC_DISPATCH))) { - return null; - } - } catch (RuntimeException e) { - caughtError = e; - } catch (Exception e) { - caughtError = new RuntimeException(e); - } - - return fireAfterExecution(hook, response, caughtError); - } - - /** - * Direct dispatch path used when no interceptor is configured. Avoids Context creation, - * hook allocation, and all hook invocations. - */ - private JsonRpcResponse handleRequestDirect( - JsonRpcRequest req, - Consumer asyncResponseCallback, - ProtocolVersion protocolVersion - ) { - try { - validate(req); - var method = req.getMethod(); - return switch (method) { - case "initialize" -> handleInitialize(req); - case "ping" -> handlePing(req); - default -> { - initializeProxies(rpcResponse -> {}); - yield switch (method) { - case "prompts/list" -> handlePromptsList(req); - case "prompts/get" -> handlePromptsGet(req); - case "tools/list" -> handleToolsList(req, protocolVersion); - case "tools/call" -> - handleToolsCallDirect(req, asyncResponseCallback, protocolVersion); - default -> methodNotFound(req); - }; - } - }; - } catch (Exception e) { - return createErrorResponse(req, e); - } - } - - private JsonRpcRequest fireBeforeExecution(McpExecutionHook hook) { - interceptor.readBeforeExecution(hook); - return interceptor.modifyBeforeExecution(hook); - } - - private JsonRpcRequest fireBeforeToolCall(McpToolCallHook hook) { - interceptor.readBeforeToolCall(hook); - return interceptor.modifyBeforeToolCall(hook); - } - - private JsonRpcResponse fireAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, - RuntimeException error - ) { - try { - interceptor.readAfterExecution(hook, response, error); - } catch (RuntimeException e) { - error = swapError("readAfterExecution", error, e); - } - try { - response = interceptor.modifyAfterExecution(hook, response, error); - error = null; - } catch (RuntimeException e) { - error = e; - } - if (error != null) { - return createErrorResponse(hook.request(), error); - } - return response; - } - - private JsonRpcResponse handleInitialize(JsonRpcRequest req) { - if (metricsObserver != null) { - var params = req.getParams(); - var clientInfo = params.getMember("clientInfo"); - var capabilities = params.getMember("capabilities"); - - String extractedProtocolVersion = params.getMember("protocolVersion") != null - ? params.getMember("protocolVersion").asString() - : null; - - String clientName = clientInfo != null && clientInfo.getMember("name") != null - ? clientInfo.getMember("name").asString() - : null; - - String clientTitle = clientInfo != null && clientInfo.getMember("title") != null - ? clientInfo.getMember("title").asString() - : null; - - boolean rootsListChanged = capabilities != null - && capabilities.getMember("roots") != null - && capabilities.getMember("roots").getMember("listChanged") != null - && capabilities.getMember("roots").getMember("listChanged").asBoolean(); - - boolean sampling = capabilities != null && capabilities.getMember("sampling") != null; - boolean elicitation = capabilities != null && capabilities.getMember("elicitation") != null; - - metricsObserver.onInitialize("initialize", - extractedProtocolVersion, - rootsListChanged, - sampling, - elicitation, - clientName, - clientTitle); - } - - this.initializeRequest.compareAndSet(null, req); - - initializeProxies(rpcResponse -> {}); - - var maybeVersion = req.getParams().getMember("protocolVersion"); - String pv = null; - if (maybeVersion != null) { - var protocolVersion = ProtocolVersion.version(maybeVersion.asString()); - if (!(protocolVersion instanceof ProtocolVersion.UnknownVersion)) { - pv = protocolVersion.identifier(); - } - } - - var builder = InitializeResult.builder(); - if (pv != null) { - builder.protocolVersion(pv); - } - - var result = builder - .capabilities(Capabilities.builder() - .tools(Tools.builder().listChanged(true).build()) - .prompts(Prompts.builder().listChanged(true).build()) - .build()) - .serverInfo(ServerInfo.builder() - .name(serviceName) - .version(version) - .build()) - .build(); - - return createSuccessResponse(req.getId(), result); - } - - private JsonRpcResponse handlePing(JsonRpcRequest req) { - return JsonRpcResponse.builder() - .id(req.getId()) - .result(Document.of(Map.of())) - .jsonrpc("2.0") - .build(); - } - - private JsonRpcResponse handlePromptsList(JsonRpcRequest req) { - var result = ListPromptsResult.builder() - .prompts(prompts.values().stream().map(Prompt::promptInfo).toList()) - .build(); - return createSuccessResponse(req.getId(), result); - } - - private JsonRpcResponse handlePromptsGet(JsonRpcRequest req) { - var promptName = req.getParams().getMember("name").asString(); - var promptArguments = req.getParams().getMember("arguments"); - - var prompt = prompts.get(normalize(promptName)); - - if (prompt == null) { - throw new RuntimeException("Prompt not found: " + promptName); - } - - var result = prompt.getPromptResult(promptArguments, req.getId()); - return createSuccessResponse(req.getId(), result); - } - - private JsonRpcResponse handleToolsList(JsonRpcRequest req, ProtocolVersion protocolVersion) { - var result = ListToolsResult.builder() - .tools(tools.values() - .stream() - .filter(t -> toolFilter.allowTool(t.serverId(), t.toolInfo().getName())) - .map(tool -> extractToolInfo(tool, protocolVersion)) - .toList()) - .build(); - return createSuccessResponse(req.getId(), result); - } - - private JsonRpcResponse handleToolsCall( - JsonRpcRequest req, - Consumer asyncResponseCallback, - ProtocolVersion protocolVersion, - McpExecutionHook executionHook - ) { - if (metricsObserver != null) { - String toolName = req.getParams().getMember("name") != null - ? req.getParams().getMember("name").asString() - : null; - metricsObserver.onToolCall("tools/call", toolName); - } - - var operationName = req.getParams().getMember("name").asString(); - var tool = tools.get(operationName); - - if (tool == null) { - return createErrorResponse(req, "No such tool: " + operationName); - } - - var toolHook = new McpToolCallHook( - req, - protocolVersion, - executionHook.context(), - operationName, - tool.serverId(), - tool.proxy() != null); - - ToolResult result; - try { - var currentReq = fireBeforeToolCall(toolHook); - toolHook = toolHook.withRequest(currentReq); - - if (tool.proxy() != null) { - return dispatchProxy(tool, currentReq, toolHook, executionHook, asyncResponseCallback); - } - - result = dispatchLocal(tool, currentReq, protocolVersion); - } catch (RuntimeException e) { - result = ToolResult.failure(e); - } - - return fireAfterToolCall(toolHook, result.response(), result.error()); - } - - private JsonRpcResponse dispatchProxy( - Tool tool, - JsonRpcRequest currentReq, - McpToolCallHook toolHook, - McpExecutionHook executionHook, - Consumer asyncResponseCallback - ) { - JsonRpcRequest proxyRequest = JsonRpcRequest.builder() - .id(currentReq.getId()) - .method(currentReq.getMethod()) - .params(currentReq.getParams()) - .jsonrpc(currentReq.getJsonrpc()) - .build(); - - executionHook.context().put(ASYNC_DISPATCH, true); - - var finalToolHook = toolHook; - tool.proxy().rpc(proxyRequest).thenAccept(response -> { - var finalResponse = fireAfterToolCall(finalToolHook, response, null); - finalResponse = fireAfterExecution(executionHook, finalResponse, null); - asyncResponseCallback.accept(finalResponse); - }).exceptionally(ex -> { - var proxyError = new RuntimeException("Proxy error: " + ex.getMessage(), ex); - var errorResponse = fireAfterToolCall(finalToolHook, null, proxyError); - if (errorResponse == null) { - errorResponse = createErrorResponse(finalToolHook.request(), proxyError); - } - errorResponse = fireAfterExecution(executionHook, errorResponse, null); - asyncResponseCallback.accept(errorResponse); - return null; - }); - - return null; - } - - private ToolResult dispatchLocal(Tool tool, JsonRpcRequest req, ProtocolVersion protocolVersion) { - try { - var operation = tool.operation(); - var argumentsDoc = req.getParams().getMember("arguments"); - var adaptedDoc = adaptDocument(argumentsDoc, operation.getApiOperation().inputSchema()); - var input = adaptedDoc.asShape(operation.getApiOperation().inputBuilder()); - var output = operation.function().apply(input, null); - var result = formatStructuredContent(tool, (SerializableShape) output, protocolVersion); - return ToolResult.success(createSuccessResponse(req.getId(), result)); - } catch (RuntimeException e) { - return ToolResult.failure(e); - } - } - - /** - * Direct tool dispatch used when no interceptor is configured. No hooks are invoked. - */ - private JsonRpcResponse handleToolsCallDirect( - JsonRpcRequest req, - Consumer asyncResponseCallback, - ProtocolVersion protocolVersion - ) { - if (metricsObserver != null) { - String toolName = req.getParams().getMember("name") != null - ? req.getParams().getMember("name").asString() - : null; - metricsObserver.onToolCall("tools/call", toolName); - } - - var operationName = req.getParams().getMember("name").asString(); - var tool = tools.get(operationName); - - if (tool == null) { - return createErrorResponse(req, "No such tool: " + operationName); - } - - if (tool.proxy() != null) { - JsonRpcRequest proxyRequest = JsonRpcRequest.builder() - .id(req.getId()) - .method(req.getMethod()) - .params(req.getParams()) - .jsonrpc(req.getJsonrpc()) - .build(); - - tool.proxy() - .rpc(proxyRequest) - .thenAccept(asyncResponseCallback) - .exceptionally(ex -> { - LOG.error("Error from proxy RPC", ex); - asyncResponseCallback.accept( - createErrorResponse(req, new RuntimeException("Proxy error: " + ex.getMessage(), ex))); - return null; - }); - return null; - } else { - var operation = tool.operation(); - var argumentsDoc = req.getParams().getMember("arguments"); - var adaptedDoc = adaptDocument(argumentsDoc, operation.getApiOperation().inputSchema()); - var input = adaptedDoc.asShape(operation.getApiOperation().inputBuilder()); - var output = operation.function().apply(input, null); - var result = formatStructuredContent(tool, (SerializableShape) output, protocolVersion); - return createSuccessResponse(req.getId(), result); - } - } - - private JsonRpcResponse fireAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, - RuntimeException error - ) { - try { - interceptor.readAfterToolCall(hook, response, error); - } catch (RuntimeException e) { - error = swapError("readAfterToolCall", error, e); - } - try { - response = interceptor.modifyAfterToolCall(hook, response, error); - error = null; - } catch (RuntimeException e) { - error = e; - } - if (error != null) { - return createErrorResponse(hook.request(), error); - } - return response; - } - - private static RuntimeException swapError(String hook, RuntimeException oldE, RuntimeException newE) { - if (oldE != null && oldE != newE) { - LOG.trace("Replacing error after {}: {} -> {}", - hook, - oldE.getClass().getName(), - newE.getClass().getName()); - } - return newE; - } - - /** - * Sets the notification writer for forwarding notifications from proxies. - */ - public void setNotificationWriter(Consumer notificationWriter) { - this.notificationWriter = notificationWriter; - } - - /** - * Creates a notification writer for a specific proxy that handles cache invalidation - * for only that proxy's tools. - */ - private Consumer createProxyNotificationWriter( - McpServerProxy proxy, - 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)); - } - } - // Forward the notification - if (baseNotificationWriter != null) { - baseNotificationWriter.accept(notification); - } - }; - } - - /** - * Starts proxies without initializing them. - */ - public void startProxies() { - for (McpServerProxy proxy : proxies.values()) { - try { - proxy.start(); - } catch (Exception e) { - LOG.error("Failed to start proxy: " + proxy.name(), e); - } - } - } - - /** - * Initializes proxies with the actual initialize request. - */ - public void initializeProxies(Consumer responseWriter) { - if (proxiesInitialized.compareAndSet(false, true)) { - JsonRpcRequest initRequest = initializeRequest.get(); - var protocolVersion = ProtocolVersion.defaultVersion(); - if (initRequest != null) { - var maybeVersion = initRequest.getParams().getMember("protocolVersion"); - if (maybeVersion != null) { - var pv = ProtocolVersion.version(maybeVersion.asString()); - if (!(pv instanceof ProtocolVersion.UnknownVersion)) { - protocolVersion = pv; - } - } - } - - for (McpServerProxy proxy : proxies.values()) { - if (initRequest != null) { - var proxyNotificationWriter = createProxyNotificationWriter(proxy, notificationWriter); - proxy.initialize(responseWriter, proxyNotificationWriter, initRequest, protocolVersion); - } - - List proxyTools = proxy.listTools(); - for (var toolInfo : proxyTools) { - tools.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); - } - } - } - } - - /** - * Gets the current initialize request if one has been received. - */ - public JsonRpcRequest getInitializeRequest() { - return initializeRequest.get(); - } - - /** - * 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))); - } - - public void addNewProxy( - McpServerProxy mcpServerProxy, - Consumer responseWriter - ) { - proxies.put(mcpServerProxy.name(), mcpServerProxy); - - 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); - } - } - - /** - * Checks if a service or proxy with the given ID exists. - */ - public boolean containsMcpServer(String id) { - return services.containsKey(id) || proxies.containsKey(id); - } - - /** - * Returns all registered proxies. - */ - public Map getProxies() { - return proxies; - } - - private boolean supportsOutputSchema(ProtocolVersion protocolVersion) { - return protocolVersion != null && protocolVersion.compareTo(ProtocolVersion.v2025_06_18.INSTANCE) >= 0; - } - - private boolean supportsAnnotations(ProtocolVersion protocolVersion) { - return protocolVersion != null && protocolVersion.compareTo(ProtocolVersion.v2025_03_26.INSTANCE) >= 0; - } - - private CallToolResult formatStructuredContent( - Tool tool, - SerializableShape output, - ProtocolVersion protocolVersion - ) { - var adaptedOutput = adaptOutputDocument(Document.of(output), tool.operation().getApiOperation().outputSchema()); - var result = CallToolResult.builder() - .content(List.of(TextContent.builder() - .text(CODEC.serializeToString(adaptedOutput)) - .build())); - - if (supportsOutputSchema(protocolVersion)) { - result.structuredContent(adaptedOutput); - } - - return result.build(); - } - - private ToolInfo extractToolInfo(Tool tool, ProtocolVersion protocolVersion) { - var toolInfo = tool.toolInfo(); - boolean stripOutput = !supportsOutputSchema(protocolVersion) && toolInfo.getOutputSchema() != null; - boolean stripAnnotations = !supportsAnnotations(protocolVersion) && toolInfo.getAnnotations() != null; - if (!stripOutput && !stripAnnotations) { - return toolInfo; - } - var builder = toolInfo.toBuilder(); - if (stripOutput) { - builder.outputSchema(null); - } - if (stripAnnotations) { - builder.annotations(null); - } - return builder.build(); - } - - private void validate(JsonRpcRequest req) { - Document id = req.getId(); - boolean isRequest = !req.getMethod().startsWith("notifications/"); - if (isRequest) { - if (id == null) { - throw ValidationException.builder() - .withoutStackTrace() - .message("Requests are expected to have ids") - .build(); - } else if (!(id.isType(ShapeType.INTEGER) || id.isType(ShapeType.STRING))) { - throw ValidationException.builder() - .withoutStackTrace() - .message("Request id is of invalid type " + id.type().name()) - .build(); - } - } - } - - private JsonRpcResponse createSuccessResponse(Document id, SerializableShape value) { - return JsonRpcResponse.builder() - .id(id) - .result(Document.of(value)) - .jsonrpc("2.0") - .build(); - } - - private JsonRpcResponse createErrorResponse(JsonRpcRequest req, Exception exception) { - return createErrorResponse(req, exception, true); //TODO change the default to false. - } - - private JsonRpcResponse createErrorResponse(JsonRpcRequest req, Throwable exception, boolean sendStackTrace) { - String s; - exception = unwrapException(exception); - if (sendStackTrace) { - try (var sw = new StringWriter(); - var pw = new PrintWriter(sw)) { - exception.printStackTrace(pw); - s = sw.toString().replace("\n", "| "); - } catch (Exception e) { - LOG.error("Error encoding response", e); - throw new RuntimeException(e); - } - } else { - s = exception.getMessage(); - } - return createErrorResponse(req, s); - } - - private Throwable unwrapException(Throwable exception) { - return switch (exception) { - case CompletionException ce when ce.getCause() != null -> ce.getCause(); - case ExecutionException ee when ee.getCause() != null -> ee.getCause(); - default -> exception; - }; - } - - private JsonRpcResponse createErrorResponse(JsonRpcRequest req, String s) { - var error = JsonRpcErrorResponse.builder() - .code(500) - .message(s) - .build(); - return JsonRpcResponse.builder() - .id(req.getId()) - .error(error) - .jsonrpc("2.0") - .build(); - } - - /** - * Per JSON-RPC 2.0, a request with an unknown method must receive a -32601 (Method not found) - * error, while notifications (requests without an id) must never receive a response. - */ - private static JsonRpcResponse methodNotFound(JsonRpcRequest req) { - if (req.getId() == null) { - return null; - } - var error = JsonRpcErrorResponse.builder() - .code(METHOD_NOT_FOUND_ERROR_CODE) - .message("Method not found: " + req.getMethod()) - .build(); - return JsonRpcResponse.builder() - .id(req.getId()) - .error(error) - .jsonrpc("2.0") - .build(); - } - - private Map createTools(Map services) { - var tools = new ConcurrentHashMap(); - for (var entry : services.entrySet()) { - var id = entry.getKey(); - var service = entry.getValue(); - var serviceName = service.schema().id().getName(); - var cache = new HashMap(); - for (var operation : service.getAllOperations()) { - var operationName = operation.name(); - Schema schema = operation.getApiOperation().schema(); - var toolInfo = ToolInfo.builder() - .name(operationName) - .description(createDescription(serviceName, - operationName, - schema)) - .inputSchema(createJsonObjectSchema( - operation.getApiOperation().inputSchema(), - operation.getApiOperation().inputSchema(), - new HashSet<>(), - cache)) - .outputSchema(createJsonObjectSchema( - operation.getApiOperation().outputSchema(), - operation.getApiOperation().outputSchema(), - new HashSet<>(), - cache)) - .annotations(createAnnotations(schema)) - .build(); - tools.put(operationName, new Tool(toolInfo, id, operation)); - } - } - return tools; - } - - private ToolAnnotations createAnnotations(Schema operationSchema) { - boolean isReadOnly = operationSchema.hasTrait(TraitKey.READ_ONLY_TRAIT); - boolean isIdempotent = operationSchema.hasTrait(TraitKey.IDEMPOTENT_TRAIT); - if (!isReadOnly && !isIdempotent) { - return null; - } - var builder = ToolAnnotations.builder(); - if (isReadOnly) { - builder.readOnlyHint(true); - } - if (isIdempotent) { - builder.idempotentHint(true); - } - return builder.build(); - } - - private JsonObjectSchema createJsonObjectSchema( - Schema member, - Schema target, - Set visited, - Map cache - ) { - var targetId = target.id(); - - var cached = cache.get(targetId); - if (cached != null) { - return (JsonObjectSchema) withDescription(cached, memberDescription(member)); - } - - if (!visited.add(targetId)) { - // if we're in a recursive cycle, just say "type": "object" and bail - return JsonObjectSchema.builder().build(); - } - - var properties = new HashMap(); - var requiredProperties = new ArrayList(); - for (var m : target.members()) { - var name = m.memberName(); - if (m.hasTrait(TraitKey.REQUIRED_TRAIT)) { - requiredProperties.add(name); - } - - var jsonSchema = createMemberSchema(m, visited, cache); - - properties.put(name, Document.of(jsonSchema)); - } - - visited.remove(targetId); - - // Cache without description so it can be reused with different member descriptions - var result = JsonObjectSchema.builder() - .properties(properties) - .required(requiredProperties) - .build(); - cache.put(targetId, result); - - return (JsonObjectSchema) withDescription(result, memberDescription(member)); - } - - private JsonArraySchema createJsonArraySchema( - Schema member, - Schema target, - Set visited, - Map cache - ) { - var listMember = target.listMember(); - var items = createMemberSchema(listMember, visited, cache); - - // For sparse lists, allow null items using anyOf - Document itemsSchema; - if (target.hasTrait(TraitKey.SPARSE_TRAIT)) { - var nullSchema = Map.of("type", Document.of("null")); - itemsSchema = Document.of(Map.of( - "anyOf", - Document.of(List.of(Document.of(items), Document.of(nullSchema))))); - } else { - itemsSchema = Document.of(items); - } - - return JsonArraySchema.builder() - .description(memberDescription(member)) - .items(itemsSchema) - .build(); - } - - private JsonPrimitiveSchema createJsonPrimitiveSchema(Schema member) { - var type = switch (member.type()) { - case BYTE, SHORT, INTEGER, INT_ENUM, LONG, FLOAT, DOUBLE -> JsonPrimitiveType.NUMBER; - case ENUM, BLOB, STRING, BIG_DECIMAL, BIG_INTEGER, TIMESTAMP -> JsonPrimitiveType.STRING; - case BOOLEAN -> JsonPrimitiveType.BOOLEAN; - default -> throw new RuntimeException(member + " is not a primitive type"); - }; - - var builder = JsonPrimitiveSchema.builder() - .type(type) - .description(memberDescription(member)); - - // Add format annotation for timestamps per JSON Schema spec - if (member.type() == ShapeType.TIMESTAMP) { - builder.format("date-time"); - } - - List enumValues = switch (member.type()) { - case ENUM, STRING -> member.stringEnumValues().stream().map(Document::of).toList(); - case INT_ENUM -> member.intEnumValues().stream().map(Document::of).toList(); - default -> List.of(); - }; - - if (!enumValues.isEmpty()) { - builder.enumValues(enumValues); - } - - return builder.build(); - } - - private static final List DOCUMENT_TYPES = List.of( - "string", - "number", - "boolean", - "object", - "array", - "null"); - - private JsonDocumentSchema createJsonDocumentSchema(Schema member) { - return JsonDocumentSchema.builder() - .type(DOCUMENT_TYPES) - .description(memberDescription(member)) - .build(); - } - - private SerializableShape createJsonDocumentSchema( - Schema member, - Set visited, - Map cache - ) { - var targetSchema = member.isMember() ? member.memberTarget() : member; - var oneOfTrait = targetSchema.getTrait(ONE_OF_TRAIT); - - if (oneOfTrait != null) { - return createJsonOneOfSchema(oneOfTrait, member, visited, cache); - } else { - return createJsonDocumentSchema(member); - } - } - - private SerializableShape createJsonOneOfSchema( - OneOfTrait oneOfTrait, - Schema documentMember, - Set visited, - Map cache - ) { - var targetId = (documentMember.isMember() ? documentMember.memberTarget() : documentMember).id(); - - var cached = cache.get(targetId); - if (cached != null) { - return withDescription(cached, memberDescription(documentMember)); - } - - if (!visited.add(targetId)) { - return JsonObjectSchema.builder().build(); - } - - var oneOfVariants = new ArrayList(); - - for (var memberDef : oneOfTrait.getMembers()) { - var memberName = memberDef.getName(); - var targetShapeId = memberDef.getTarget(); - - var targetSchema = schemaIndex.getSchema(targetShapeId); - var memberSchema = createJsonObjectSchema(targetSchema, targetSchema, visited, cache); - - oneOfVariants.add(createUnionVariant(memberName, memberSchema)); - } - - visited.remove(targetId); - - var result = JsonOneOfSchema.builder() - .oneOf(oneOfVariants) - .build(); - cache.put(targetId, result); - - return withDescription(result, memberDescription(documentMember)); - } - - private SerializableShape createJsonUnionSchema( - Schema member, - Schema target, - Set visited, - Map cache - ) { - var targetId = target.id(); - - var cached = cache.get(targetId); - if (cached != null) { - return withDescription(cached, memberDescription(member)); - } - - if (!visited.add(targetId)) { - return JsonObjectSchema.builder().build(); - } - - var variants = new ArrayList(); - - for (var m : target.members()) { - var memberName = m.memberName(); - var memberSchema = createMemberSchema(m, visited, cache); - - variants.add(createUnionVariant(memberName, memberSchema)); - } - - visited.remove(targetId); - - var result = JsonOneOfSchema.builder() - .oneOf(variants) - .build(); - cache.put(targetId, result); - - return withDescription(result, memberDescription(member)); - } - - private static Document createUnionVariant(String memberName, SerializableShape memberSchema) { - var wrapperSchema = JsonObjectSchema.builder() - .properties(Map.of(memberName, Document.of(memberSchema))) - .required(List.of(memberName)) - .additionalProperties(Document.of(false)) - .build(); - return Document.of(wrapperSchema); - } - - private SerializableShape createMemberSchema( - Schema member, - Set visited, - Map cache - ) { - return switch (member.type()) { - case LIST, SET -> createJsonArraySchema(member, member.memberTarget(), visited, cache); - case MAP -> createJsonMapSchema(member, member.memberTarget(), visited, cache); - case STRUCTURE -> createJsonObjectSchema(member, member.memberTarget(), visited, cache); - case UNION -> createJsonUnionSchema(member, member.memberTarget(), visited, cache); - case DOCUMENT -> createJsonDocumentSchema(member, visited, cache); - default -> createJsonPrimitiveSchema(member); - }; - } - - private JsonObjectSchema createJsonMapSchema( - Schema member, - Schema target, - Set visited, - Map cache - ) { - var mapValueMember = target.mapValueMember(); - var valueSchema = createMemberSchema(mapValueMember, visited, cache); - - // For sparse maps, allow null values using anyOf - Document additionalPropertiesSchema; - if (target.hasTrait(TraitKey.SPARSE_TRAIT)) { - var nullSchema = Map.of("type", Document.of("null")); - additionalPropertiesSchema = Document.of(Map.of( - "anyOf", - Document.of(List.of(Document.of(valueSchema), Document.of(nullSchema))))); - } else { - additionalPropertiesSchema = Document.of(valueSchema); - } - - return JsonObjectSchema.builder() - .description(memberDescription(member)) - .additionalProperties(additionalPropertiesSchema) - .build(); - } - - private static String memberDescription(Schema schema) { - String description = null; - // Use getDirectTrait for members to avoid inheriting the target's documentation trait - // (getTrait on a member merges member + target traits, which would cause doubling) - var trait = schema.isMember() - ? schema.getDirectTrait(TraitKey.DOCUMENTATION_TRAIT) - : schema.getTrait(TraitKey.DOCUMENTATION_TRAIT); - if (trait != null) { - description = trait.getValue(); - } - if (schema.isMember()) { - var memberDescription = memberDescription(schema.memberTarget()); - if (description != null && memberDescription != null) { - description = appendSentences(description, memberDescription); - } else if (memberDescription != null) { - description = memberDescription; - } - } - return description; - } - - private static String createDescription( - String serviceName, - String operationName, - Schema schema - ) { - var documentationTrait = schema.getTrait(TraitKey.DOCUMENTATION_TRAIT); - if (documentationTrait != null) { - return documentationTrait.getValue(); - } else { - return "This tool invokes %s API of %s.".formatted(operationName, serviceName); - } - } - - private record Tool( - ToolInfo toolInfo, - String serverId, - Operation operation, - McpServerProxy proxy, - boolean requiredAdapting) { - - Tool(ToolInfo toolInfo, String serverId, Operation operation) { - this(toolInfo, serverId, operation, null, false); - } - - Tool(ToolInfo toolInfo, String serverId, McpServerProxy proxy) { - this(toolInfo, serverId, null, proxy, false); - } - } - - private record ToolResult(JsonRpcResponse response, RuntimeException error) { - static ToolResult success(JsonRpcResponse response) { - return new ToolResult(response, null); - } - - static ToolResult failure(RuntimeException error) { - return new ToolResult(null, error); - } - } - - private static String appendSentences(String first, String second) { - first = first.trim(); - if (!first.endsWith(".")) { - first = first + ". "; - } - return first + second; - } - - private static SerializableShape withDescription(SerializableShape schema, String description) { - if (description == null) { - return schema; - } - if (schema instanceof JsonObjectSchema s) { - return s.toBuilder().description(description).build(); - } - if (schema instanceof JsonOneOfSchema s) { - return s.toBuilder().description(description).build(); - } - return schema; - } - - private Document adaptDocument(Document doc, Schema schema) { - if (doc == null) { - return null; - } - var fromType = doc.type(); - var toType = schema.type(); - return switch (toType) { - case BIG_DECIMAL -> switch (fromType) { - case STRING -> Document.of(new BigDecimal(doc.asString())); - case BIG_INTEGER -> doc; - default -> badType(fromType, toType); - }; - case BIG_INTEGER -> switch (fromType) { - case STRING -> Document.of(new BigInteger(doc.asString())); - case BIG_INTEGER -> doc; - default -> badType(fromType, toType); - }; - case BLOB -> switch (fromType) { - case STRING -> Document.of(Base64.getDecoder().decode(doc.asString())); - case BLOB -> doc; - default -> badType(fromType, toType); - }; - case TIMESTAMP -> adaptTimestamp(doc); - case STRUCTURE -> { - var convertedMembers = new HashMap(); - var members = schema.members(); - for (var member : members) { - var memberName = member.memberName(); - var memberDoc = doc.getMember(memberName); - if (memberDoc != null) { - convertedMembers.put(memberName, adaptDocument(memberDoc, member)); - } - } - yield Document.of(convertedMembers); - } - case UNION -> { - var convertedMembers = new HashMap(); - - // Find which member is set and adapt it - // Input is in wrapper format: {"circle": {...}} - for (var member : schema.members()) { - var memberName = member.memberName(); - var memberDoc = doc.getMember(memberName); - if (memberDoc != null) { - convertedMembers.put(memberName, adaptDocument(memberDoc, member)); - break; - } - } - yield Document.of(convertedMembers); - } - case LIST, SET -> { - var listMember = schema.listMember(); - var convertedList = new ArrayList(); - for (var item : doc.asList()) { - convertedList.add(adaptDocument(item, listMember)); - } - yield Document.of(convertedList); - } - case MAP -> { - var mapValue = schema.mapValueMember(); - var convertedMap = new HashMap(); - for (var entry : doc.asStringMap().entrySet()) { - convertedMap.put(entry.getKey(), adaptDocument(entry.getValue(), mapValue)); - } - yield Document.of(convertedMap); - } - case DOCUMENT -> adaptDocumentWithOneOf(doc, schema); - default -> doc; - }; - } - - private Document adaptDocumentWithOneOf(Document doc, Schema schema) { - var targetSchema = schema.isMember() ? schema.memberTarget() : schema; - var oneOfTrait = targetSchema.getTrait(ONE_OF_TRAIT); - - if (oneOfTrait != null) { - // MCP sends wrapper format: {"circle": {"radius": 5}} - // Need to convert to discriminated format: {"__type": "smithy.example#Circle", "radius": 5} - var discriminator = oneOfTrait.getDiscriminator(); - - // Find which member is set in the wrapper - for (var memberDef : oneOfTrait.getMembers()) { - var memberName = memberDef.getName(); - var memberDoc = doc.getMember(memberName); - if (memberDoc != null) { - // Build the flat object with discriminator - var flatMembers = new HashMap(); - var memberId = memberDef.getTarget(); - flatMembers.put(discriminator, Document.of(memberId.toString())); - // Copy all fields from the inner object - var memberSchema = schemaIndex.getSchema(memberId); - flatMembers.putAll(adaptDocument(memberDoc, memberSchema).asStringMap()); - return Document.of(flatMembers); - } - } - // Fallback - return as-is if can't determine type - } - return doc; - } - - private static Document badType(ShapeType from, ShapeType to) { - throw new RuntimeException("Cannot convert from " + from + " to " + to); - } - - /** - * This is primarily for more robustness against AI hallucinations. - */ - private static Document adaptTimestamp(Document doc) { - // If already a timestamp, format as date-time string - if (doc.isType(ShapeType.TIMESTAMP)) { - return Document.of(DATE_TIME.writeString(doc.asTimestamp())); - } - // If input is a string, try DATE_TIME first, fallback to HTTP_DATE - if (doc.isType(ShapeType.STRING)) { - var str = doc.asString(); - try { - return Document.of(DATE_TIME.readFromString(str, false)); - } catch (Exception e) { - // Fallback to HTTP_DATE format - return Document.of(HTTP_DATE.readFromString(str, false)); - } - } - // If input is a number, use epoch seconds - return Document.of(EPOCH_SECONDS.readFromNumber(doc.asNumber())); - } - - private Document adaptOutputDocument(Document doc, Schema schema) { - if (doc == null) { - return null; - } - var toType = schema.type(); - return switch (toType) { - case BIG_DECIMAL -> Document.of(doc.asBigDecimal().toString()); - case BIG_INTEGER -> Document.of(doc.asBigInteger().toString()); - case BLOB -> Document.of(Base64.getEncoder().encodeToString(ByteBufferUtils.getBytes(doc.asBlob()))); - // Use adaptTimestamp() instead of asTimestamp() because oneOf union members are - // deserialized as untyped Documents (no schema available). Timestamps in these - // documents remain as strings or numbers rather than being converted to Timestamp Documents. - case TIMESTAMP -> adaptTimestamp(doc); - case STRUCTURE -> { - var convertedMembers = new HashMap(); - for (var member : schema.members()) { - var memberName = member.memberName(); - var memberDoc = doc.getMember(memberName); - if (memberDoc != null) { - convertedMembers.put(memberName, adaptOutputDocument(memberDoc, member)); - } - } - yield Document.of(convertedMembers); - } - case UNION -> { - // Regular union - already in wrapper format: {"circle": {...}} - for (var member : schema.members()) { - var memberName = member.memberName(); - var memberDoc = doc.getMember(memberName); - if (memberDoc != null) { - var adaptedMemberDoc = adaptOutputDocument(memberDoc, member); - yield Document.of(Map.of(memberName, adaptedMemberDoc)); - } - } - yield Document.of(Map.of()); - } - case LIST, SET -> { - var listMember = schema.listMember(); - var convertedList = new ArrayList(); - for (var item : doc.asList()) { - convertedList.add(adaptOutputDocument(item, listMember)); - } - yield Document.of(convertedList); - } - case MAP -> { - var mapValue = schema.mapValueMember(); - var convertedMap = new HashMap(); - for (var entry : doc.asStringMap().entrySet()) { - convertedMap.put(entry.getKey(), adaptOutputDocument(entry.getValue(), mapValue)); - } - yield Document.of(convertedMap); - } - case DOCUMENT -> { - var targetSchema = schema.isMember() ? schema.memberTarget() : schema; - var oneOfTrait = targetSchema.getTrait(ONE_OF_TRAIT); - - if (oneOfTrait != null) { - // External service returns: {"__type": "smithy.example#Circle", "radius": 5} - // Need to convert to MCP wrapper format: {"circle": {"radius": 5}} - var discriminator = oneOfTrait.getDiscriminator(); - var discriminatorValue = doc.getMember(discriminator); - - if (discriminatorValue != null) { - var shapeId = ShapeId.from(discriminatorValue.asString()); - // Find the matching member definition - for (var memberDef : oneOfTrait.getMembers()) { - if (memberDef.getTarget().equals(shapeId)) { - var memberName = memberDef.getName(); - var memberSchema = schemaIndex.getSchema(shapeId); - // Build the inner object without the discriminator field - var innerMembers = new HashMap<>(adaptOutputDocument(doc, memberSchema).asStringMap()); - innerMembers.remove(discriminator); - // Return wrapper format - yield Document.of(Map.of(memberName, Document.of(innerMembers))); - } - } - } - } - yield doc; - } - default -> doc; - }; - } - - public static Builder builder() { - return new Builder(); - } - - public static class Builder { - private Map services = new HashMap<>(); - private List proxyList = new ArrayList<>(); - private McpServerInterceptor interceptor = McpServerInterceptor.NOOP; - private String name = "mcp-server"; - private String version = "1.0.0"; - private ToolFilter toolFilter = (serverId, toolName) -> true; - private McpMetricsObserver metricsObserver; - - public Builder services(Map services) { - this.services = services; - return this; - } - - public Builder proxyList(List proxyList) { - this.proxyList = proxyList; - return this; - } - - public Builder name(String name) { - this.name = name; - return this; - } - - public Builder version(String version) { - this.version = version; - return this; - } - - public Builder toolFilter(ToolFilter toolFilter) { - this.toolFilter = toolFilter; - return this; - } - - public Builder metricsObserver(McpMetricsObserver metricsObserver) { - this.metricsObserver = metricsObserver; - return this; - } - - /** - * Sets the server interceptor. Use {@link McpServerInterceptor#chain(List)} to compose - * multiple interceptors into one. - * - * @see McpServerInterceptor for hook descriptions and the execution lifecycle - */ - public Builder interceptor(McpServerInterceptor interceptor) { - this.interceptor = Objects.requireNonNull(interceptor, "interceptor"); - return this; - } - - public McpService build() { - return new McpService(services, proxyList, name, version, toolFilter, metricsObserver, interceptor); - } - } -} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSession.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSession.java new file mode 100644 index 0000000000..cfe0b6e71c --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSession.java @@ -0,0 +1,44 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.concurrent.locks.ReentrantLock; + +final class McpSession { + private final ReentrantLock negotiationLock = new ReentrantLock(); + private final McpProtocolRegistry protocols; + private ProtocolVersion version; + + McpSession(McpProtocolRegistry protocols) { + this.protocols = protocols; + version = protocols.defaultProtocol().protocolVersion(); + } + + ProtocolVersion negotiate(McpCall call, ProtocolVersion transportClaim) { + negotiationLock.lock(); + try { + var claimed = call.metadata().protocolVersion(); + if (claimed != null) { + version = protocols.require(claimed).protocolVersion(); + return version; + } + if (call instanceof McpCall.Initialize initialize) { + var requested = initialize.requestedVersion(); + var requestedProtocol = protocols.find(requested); + version = requestedProtocol != null && !requestedProtocol.usesStatelessMetadata() + ? requestedProtocol.protocolVersion() + : protocols.initializationFallbackProtocol().protocolVersion(); + return version; + } + if (transportClaim != null) { + version = protocols.require(transportClaim).protocolVersion(); + } + return version; + } finally { + negotiationLock.unlock(); + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSourceSnapshot.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSourceSnapshot.java new file mode 100644 index 0000000000..053c0cf96b --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSourceSnapshot.java @@ -0,0 +1,13 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Map; + +record McpSourceSnapshot( + Map tools, + Map prompts, + SmithyDocumentAdapter documentAdapter) {} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSources.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSources.java new file mode 100644 index 0000000000..5a30896b0e --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSources.java @@ -0,0 +1,45 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Map; +import java.util.function.Consumer; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.java.mcp.model.JsonRpcResponse; +import software.amazon.smithy.java.server.Service; + +/** + * Aggregates local services and remote MCP peers behind one immutable-snapshot source. + */ +interface McpSources extends AutoCloseable { + McpSourceSnapshot snapshot(); + + McpToolDescriptor tool(String name); + + McpPromptDescriptor prompt(String normalizedName); + + Map remoteClients(); + + boolean containsServer(String id); + + void bindTransport( + Consumer notificationWriter, + Consumer responseWriter + ); + + void initializeRemoteClients(JsonRpcRequest request, McpProtocol protocol); + + void ensureRemoteCatalogLoaded(); + + void addService(String id, Service service); + + void addRemoteClient(McpRemoteClient client); + + Map headerParameters(String toolName); + + @Override + void close(); +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolCallHook.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolCallHook.java deleted file mode 100644 index cff5813afd..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolCallHook.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import software.amazon.smithy.java.context.Context; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; -import software.amazon.smithy.utils.SmithyUnstableApi; - -/** - * Hook data available during tool call processing. Extends {@link McpExecutionHook} with - * tool-specific information. Passed to tool-scoped hooks in {@link McpServerInterceptor}. - */ -@SmithyUnstableApi -public class McpToolCallHook extends McpExecutionHook { - - private final String toolName; - private final String serverId; - private final boolean isProxy; - - McpToolCallHook( - JsonRpcRequest request, - ProtocolVersion protocolVersion, - Context context, - String toolName, - String serverId, - boolean isProxy - ) { - super(request, protocolVersion, context); - this.toolName = toolName; - this.serverId = serverId; - this.isProxy = isProxy; - } - - /** - * The name of the tool being invoked. - */ - public String toolName() { - return toolName; - } - - /** - * The server ID that owns this tool. - */ - public String serverId() { - return serverId; - } - - /** - * Whether this tool is dispatched to a remote proxy rather than handled locally. - */ - public boolean isProxy() { - return isProxy; - } - - /** - * Returns a new hook with the given request, or the same hook if unchanged. - */ - @Override - public McpToolCallHook withRequest(JsonRpcRequest request) { - return this.request() == request - ? this - : new McpToolCallHook(request, protocolVersion(), context(), toolName, serverId, isProxy); - } -} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolDescriptor.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolDescriptor.java new file mode 100644 index 0000000000..cdd376db6d --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolDescriptor.java @@ -0,0 +1,26 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Map; +import software.amazon.smithy.java.mcp.model.ToolInfo; +import software.amazon.smithy.java.server.Operation; + +record McpToolDescriptor( + ToolInfo info, + String serverId, + Target target, + Map headerParameters) { + McpToolDescriptor { + headerParameters = Map.copyOf(headerParameters); + } + + sealed interface Target permits LocalTarget, RemoteTarget {} + + record LocalTarget(Operation operation) implements Target {} + + record RemoteTarget(McpRemoteClient client) implements Target {} +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutionContext.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutionContext.java new file mode 100644 index 0000000000..eddb351e35 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutionContext.java @@ -0,0 +1,29 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Typed data exposed around tool execution. + */ +@SmithyUnstableApi +public record McpToolExecutionContext( + McpCall.CallTool call, + McpRequestContext requestContext, + String serverId, + boolean remote) { + public McpToolExecutionContext { + Objects.requireNonNull(call, "call"); + Objects.requireNonNull(requestContext, "requestContext"); + Objects.requireNonNull(serverId, "serverId"); + } + + McpToolExecutionContext withCall(McpCall.CallTool call) { + return this.call == call ? this : new McpToolExecutionContext(call, requestContext, serverId, remote); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutor.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutor.java new file mode 100644 index 0000000000..572881816d --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutor.java @@ -0,0 +1,145 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import software.amazon.smithy.java.core.schema.SerializableShape; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.CallToolResult; +import software.amazon.smithy.java.mcp.model.TextContent; + +/** + * Executes local and remote tool targets behind one blocking interface. + */ +final class McpToolExecutor { + private final McpSources sources; + private final McpWireCodec wireCodec; + private final McpInterceptor interceptor; + private final McpProtocolRegistry protocols; + + McpToolExecutor( + McpSources sources, + McpWireCodec wireCodec, + McpInterceptor interceptor, + McpProtocolRegistry protocols + ) { + this.sources = sources; + this.wireCodec = wireCodec; + this.interceptor = interceptor; + this.protocols = protocols; + } + + McpOutcome execute(McpCall.CallTool call, McpRequestContext requestContext) { + var descriptor = sources.tool(call.name()); + if (descriptor == null) { + return new McpOutcome.Failure( + call.id(), + new McpError(-32602, "No such tool: " + call.name(), null)); + } + + var toolContext = new McpToolExecutionContext( + call, + requestContext, + descriptor.serverId(), + descriptor.target() instanceof McpToolDescriptor.RemoteTarget); + McpOutcome outcome = null; + RuntimeException error = null; + try { + interceptor.readBeforeToolCall(toolContext); + call = interceptor.modifyBeforeToolCall(toolContext); + toolContext = toolContext.withCall(call); + outcome = invoke(descriptor, call, requestContext); + } catch (RuntimeException e) { + error = e; + } + + try { + interceptor.readAfterToolCall(toolContext, outcome, error); + } catch (RuntimeException e) { + if (error == null) { + error = e; + } else if (error != e) { + error.addSuppressed(e); + } + } + return interceptor.modifyAfterToolCall(toolContext, outcome, error); + } + + private McpOutcome invoke( + McpToolDescriptor descriptor, + McpCall.CallTool call, + McpRequestContext requestContext + ) { + return switch (descriptor.target()) { + case McpToolDescriptor.LocalTarget local -> invokeLocal(descriptor, local, call, requestContext); + case McpToolDescriptor.RemoteTarget remote -> invokeRemote(remote, call); + }; + } + + private McpOutcome invokeLocal( + McpToolDescriptor descriptor, + McpToolDescriptor.LocalTarget target, + McpCall.CallTool call, + McpRequestContext requestContext + ) { + var operation = target.operation(); + var adapter = sources.snapshot().documentAdapter(); + var inputDocument = adapter.toSmithy(call.arguments(), operation.getApiOperation().inputSchema()); + var input = inputDocument.asShape(operation.getApiOperation().inputBuilder()); + + final SerializableShape output; + try { + output = (SerializableShape) operation.function().apply(input, null); + } catch (RuntimeException e) { + return toolFailure(call, e); + } + + var outputDocument = adapter.fromSmithy( + Document.of(output), + operation.getApiOperation().outputSchema()); + var result = CallToolResult.builder() + .content(List.of(TextContent.builder() + .text(McpJson.CODEC.serializeToString(outputDocument)) + .build())); + var protocol = protocols.require(requestContext.protocolVersion()); + if (protocol.supportsOutputSchema()) { + result.structuredContent(outputDocument); + } + return new McpOutcome.Success(call.id(), Document.of(result.build())); + } + + private McpOutcome invokeRemote( + McpToolDescriptor.RemoteTarget target, + McpCall.CallTool call + ) { + return wireCodec.decode(target.client().exchange(wireCodec.encode(call))); + } + + private McpOutcome toolFailure(McpCall.CallTool call, RuntimeException exception) { + var cause = unwrap(exception); + var message = cause.getMessage(); + if (message == null || message.isBlank()) { + message = cause.getClass().getSimpleName(); + } + var result = CallToolResult.builder() + .content(List.of(TextContent.builder().text(message).build())) + .isError(true) + .build(); + return new McpOutcome.Success(call.id(), Document.of(result)); + } + + private Throwable unwrap(Throwable exception) { + return switch (exception) { + case CompletionException completion when completion.getCause() != null -> + completion.getCause(); + case ExecutionException execution when execution.getCause() != null -> + execution.getCause(); + default -> exception; + }; + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpTransportContext.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpTransportContext.java new file mode 100644 index 0000000000..c141bcfebb --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpTransportContext.java @@ -0,0 +1,26 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.List; +import java.util.Map; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Transport-specific request information available during protocol validation. + */ +@SmithyUnstableApi +public interface McpTransportContext { + McpTransportContext STDIO = new Stdio(); + + record Stdio() implements McpTransportContext {} + + record Http(Map> headers, boolean loopbackOnly) implements McpTransportContext { + public Http { + headers = headers == null ? Map.of() : Map.copyOf(headers); + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpUnsupportedMethodException.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpUnsupportedMethodException.java new file mode 100644 index 0000000000..5a2a522a14 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpUnsupportedMethodException.java @@ -0,0 +1,15 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +/** + * Internal control signal used only when a selected protocol does not define a method. + */ +final class McpUnsupportedMethodException extends RuntimeException { + McpUnsupportedMethodException(McpMethod method, McpProtocolId protocolId) { + super(method.wireName() + " is not supported by MCP " + protocolId.identifier()); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireCodec.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireCodec.java new file mode 100644 index 0000000000..a408094028 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireCodec.java @@ -0,0 +1,166 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.HashMap; +import java.util.Map; +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; + +final class McpWireCodec { + private final McpRequestDecoder decoder; + + McpWireCodec(Map> extensions) { + decoder = new McpRequestDecoder(extensions); + } + + McpCall decode(JsonRpcRequest request) { + return decoder.decode(request); + } + + JsonRpcRequest encode(McpCall call) { + var params = switch (call) { + case McpCall.Initialize c -> initializeParams(c); + case McpCall.Ping ignored -> Document.of(Map.of()); + case McpCall.Discover ignored -> Document.of(Map.of()); + case McpCall.ListTools c -> optionalParam("cursor", c.cursor()); + case McpCall.CallTool c -> Document.of(Map.of( + "name", + Document.of(c.name()), + "arguments", + c.arguments())); + case McpCall.ListPrompts c -> optionalParam("cursor", c.cursor()); + case McpCall.GetPrompt c -> { + var values = new HashMap(); + values.put("name", Document.of(c.name())); + if (!c.arguments().isEmpty()) { + values.put("arguments", Document.of(c.arguments())); + } + yield Document.of(values); + } + case McpCall.Complete c -> completionParams(c); + case McpCall.SetLogLevel c -> Document.of(Map.of("level", Document.of(c.level()))); + case McpCall.ReadResource c -> Document.of(Map.of("uri", Document.of(c.uri()))); + case McpCall.Notification c -> c.params(); + case McpCall.ExtensionCall extension -> encodeExtension(extension); + case McpCall.UnknownCall c -> c.params(); + }; + params = withMetadata(params, call.metadata()); + return JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(call.id()) + .method(call.method().wireName()) + .params(params) + .build(); + } + + JsonRpcResponse encode(McpOutcome outcome) { + return switch (outcome) { + case McpOutcome.Success success -> JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(success.id()) + .result(success.result()) + .build(); + case McpOutcome.Failure failure -> { + var error = JsonRpcErrorResponse.builder() + .code(failure.error().code()) + .message(failure.error().message()); + if (failure.error().data() != null) { + error.data(failure.error().data()); + } + yield JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(failure.id()) + .error(error.build()) + .build(); + } + case McpOutcome.NoResponse ignored -> null; + }; + } + + McpOutcome decode(JsonRpcResponse response) { + if (response == null) { + return McpOutcome.NoResponse.INSTANCE; + } + if (response.getError() != null) { + return new McpOutcome.Failure( + response.getId(), + new McpError( + response.getError().getCode(), + response.getError().getMessage(), + response.getError().getData())); + } + return new McpOutcome.Success(response.getId(), response.getResult()); + } + + private

Document encodeExtension(McpCall.ExtensionCall

call) { + return call.extension().encode(call.parameters()); + } + + private Document initializeParams(McpCall.Initialize call) { + var values = new HashMap(); + values.put("protocolVersion", Document.of(call.requestedVersion().identifier())); + if (call.clientInfo() != null) { + values.put("clientInfo", call.clientInfo()); + } + if (call.capabilities() != null) { + values.put("capabilities", call.capabilities()); + } + return Document.of(values); + } + + private Document completionParams(McpCall.Complete call) { + var values = new HashMap(); + if (call.reference() != null) { + var reference = new HashMap(); + if (call.reference().type() != null) { + reference.put("type", Document.of(call.reference().type())); + } + if (call.reference().name() != null) { + reference.put("name", Document.of(call.reference().name())); + } + values.put("ref", Document.of(reference)); + } + if (call.argument() != null) { + var argument = new HashMap(); + if (call.argument().name() != null) { + argument.put("name", Document.of(call.argument().name())); + } + if (call.argument().value() != null) { + argument.put("value", Document.of(call.argument().value())); + } + values.put("argument", Document.of(argument)); + } + return Document.of(values); + } + + private Document optionalParam(String name, String value) { + return value == null ? Document.of(Map.of()) : Document.of(Map.of(name, Document.of(value))); + } + + private Document withMetadata(Document params, McpMetadata metadata) { + if (metadata == null || metadata == McpMetadata.EMPTY) { + return params; + } + var values = new HashMap<>(params.asStringMap()); + var meta = new HashMap<>(metadata.extensions()); + if (metadata.protocolVersion() != null) { + meta.put(McpWireNames.PROTOCOL_VERSION, Document.of(metadata.protocolVersion().identifier())); + } + if (metadata.clientInfo() != null) { + meta.put(McpWireNames.CLIENT_INFO, metadata.clientInfo()); + } + if (metadata.clientCapabilities() != null) { + meta.put(McpWireNames.CLIENT_CAPABILITIES, metadata.clientCapabilities()); + } + if (!meta.isEmpty()) { + values.put("_meta", Document.of(meta)); + } + return Document.of(values); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireNames.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireNames.java new file mode 100644 index 0000000000..c30b2c10d3 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireNames.java @@ -0,0 +1,15 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +final class McpWireNames { + static final String PROTOCOL_VERSION = "io.modelcontextprotocol/protocolVersion"; + static final String CLIENT_INFO = "io.modelcontextprotocol/clientInfo"; + static final String CLIENT_CAPABILITIES = "io.modelcontextprotocol/clientCapabilities"; + static final String SERVER_INFO = "io.modelcontextprotocol/serverInfo"; + + private McpWireNames() {} +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/Prompt.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/Prompt.java index 51e54ddede..ad45e3c725 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/Prompt.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/Prompt.java @@ -33,7 +33,7 @@ public final class Prompt { private final PromptInfo promptInfo; private final String promptTemplate; - private final McpServerProxy proxy; + private final McpRemoteClient proxy; /** * Creates a local prompt with a template. @@ -53,7 +53,7 @@ public Prompt(PromptInfo promptInfo, String promptTemplate) { * @param promptInfo The prompt metadata * @param proxy The MCP server proxy to delegate to */ - public Prompt(PromptInfo promptInfo, McpServerProxy proxy) { + public Prompt(PromptInfo promptInfo, McpRemoteClient proxy) { this.promptInfo = promptInfo; this.promptTemplate = null; this.proxy = proxy; @@ -98,12 +98,11 @@ private GetPromptResult delegateToProxy(Document arguments, Document requestId) .jsonrpc("2.0") .build(); - return proxy.rpc(request).thenApply(response -> { - if (response.getError() != null) { - throw new RuntimeException("Error getting prompt: " + response.getError().getMessage()); - } - return response.getResult().asShape(GetPromptResult.builder()); - }).join(); + var response = proxy.exchange(request); + if (response.getError() != null) { + throw new McpRemoteException("Error getting prompt: " + response.getError().getMessage()); + } + return response.getResult().asShape(GetPromptResult.builder()); } /** diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ProtocolVersion.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ProtocolVersion.java index 555559a3b8..3b381a7c3c 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ProtocolVersion.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ProtocolVersion.java @@ -7,82 +7,37 @@ import software.amazon.smithy.utils.SmithyUnstableApi; +/** + * A protocol version claimed by an MCP peer. + * + *

Built-in versions are represented by {@link KnownProtocolVersion}; other wire + * values are preserved in {@link UnknownProtocolVersion} and may resolve to a + * registered {@link ExtensionMcpProtocol}. + */ @SmithyUnstableApi -public abstract sealed class ProtocolVersion implements Comparable - permits ProtocolVersion.UnknownVersion, ProtocolVersion.v2024_11_05, ProtocolVersion.v2025_03_26, - ProtocolVersion.v2025_06_18, ProtocolVersion.v2025_11_25 { - public static final class v2025_11_25 extends ProtocolVersion { - public static final v2025_11_25 INSTANCE = new v2025_11_25(); - - private v2025_11_25() { - super("2025-11-25"); - } - } - - public static final class v2025_06_18 extends ProtocolVersion { - public static final v2025_06_18 INSTANCE = new v2025_06_18(); - - private v2025_06_18() { - super("2025-06-18"); - } - } - - public static final class v2025_03_26 extends ProtocolVersion { - public static final v2025_03_26 INSTANCE = new v2025_03_26(); - - private v2025_03_26() { - super("2025-03-26"); - } - } - - public static final class v2024_11_05 extends ProtocolVersion { - public static final v2024_11_05 INSTANCE = new v2024_11_05(); - - private v2024_11_05() { - super("2024-11-05"); +public sealed interface ProtocolVersion permits KnownProtocolVersion, UnknownProtocolVersion { + + /** + * Returns the wire identifier of this version. + */ + String identifier(); + + /** + * Parses a wire protocol version. + */ + static ProtocolVersion parse(String identifier) { + if (identifier == null) { + return defaultVersion(); } + var known = KnownProtocolVersion.fromIdentifier(identifier); + return known == null ? new UnknownProtocolVersion(identifier) : known; } - public static final class UnknownVersion extends ProtocolVersion { - private UnknownVersion(String identifier) { - super(identifier); - } - } - - private final String identifier; - - private ProtocolVersion(String identifier) { - this.identifier = identifier; - } - - public String identifier() { - return identifier; - } - - @Override - public final int compareTo(ProtocolVersion o) { - if (o instanceof UnknownVersion) { - if (this instanceof UnknownVersion) { - return 0; - } - return 1; - } - - return identifier.compareTo(o.identifier); - } - - public static ProtocolVersion version(String identifier) { - return switch (identifier) { - case null -> v2025_03_26.INSTANCE; - case "2024-11-05" -> v2024_11_05.INSTANCE; - case "2025-03-26" -> v2025_03_26.INSTANCE; - case "2025-06-18" -> v2025_06_18.INSTANCE; - case "2025-11-25" -> v2025_11_25.INSTANCE; - default -> new UnknownVersion(identifier); - }; - } - - public static ProtocolVersion defaultVersion() { - return v2025_03_26.INSTANCE; + /** + * The compatibility version used when a pre-2025-06-18 HTTP client omits the + * protocol-version header. + */ + static KnownProtocolVersion defaultVersion() { + return KnownProtocolVersion.V2025_03_26; } } diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/SmithyDocumentAdapter.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/SmithyDocumentAdapter.java new file mode 100644 index 0000000000..2e77dc891b --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/SmithyDocumentAdapter.java @@ -0,0 +1,253 @@ +/* + * 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 software.amazon.smithy.java.core.serde.TimestampFormatter.Prelude.DATE_TIME; +import static software.amazon.smithy.java.core.serde.TimestampFormatter.Prelude.EPOCH_SECONDS; +import static software.amazon.smithy.java.core.serde.TimestampFormatter.Prelude.HTTP_DATE; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.HashMap; +import java.util.IdentityHashMap; +import java.util.Map; +import java.util.Set; +import software.amazon.smithy.java.core.schema.Schema; +import software.amazon.smithy.java.core.schema.SchemaIndex; +import software.amazon.smithy.java.core.schema.TraitKey; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.mcp.OneOfTrait; +import software.amazon.smithy.model.shapes.ShapeId; +import software.amazon.smithy.model.shapes.ShapeType; + +/** + * Adapts schemaless MCP documents to and from Smithy runtime values. + */ +final class SmithyDocumentAdapter { + private static final TraitKey ONE_OF_TRAIT = TraitKey.get(OneOfTrait.class); + + private final SchemaIndex schemaIndex; + private final Map adaptationRequired = + Collections.synchronizedMap(new IdentityHashMap<>()); + + SmithyDocumentAdapter(SchemaIndex schemaIndex) { + this.schemaIndex = schemaIndex; + } + + Document toSmithy(Document document, Schema schema) { + if (document == null) { + return null; + } + if (!needsAdaptation(schema)) { + return document; + } + var fromType = document.type(); + var toType = schema.type(); + return switch (toType) { + case BIG_DECIMAL -> switch (fromType) { + case STRING -> Document.of(new BigDecimal(document.asString())); + case BIG_INTEGER -> document; + default -> badType(fromType, toType); + }; + case BIG_INTEGER -> switch (fromType) { + case STRING -> Document.of(new BigInteger(document.asString())); + case BIG_INTEGER -> document; + default -> badType(fromType, toType); + }; + case BLOB -> switch (fromType) { + case STRING -> Document.of(Base64.getDecoder().decode(document.asString())); + case BLOB -> document; + default -> badType(fromType, toType); + }; + case TIMESTAMP -> adaptTimestamp(document); + case STRUCTURE -> { + var converted = new HashMap(); + for (var member : schema.members()) { + var memberDocument = document.getMember(member.memberName()); + if (memberDocument != null) { + converted.put(member.memberName(), toSmithy(memberDocument, member)); + } + } + yield Document.of(converted); + } + case UNION -> { + var converted = new HashMap(); + for (var member : schema.members()) { + var memberDocument = document.getMember(member.memberName()); + if (memberDocument != null) { + converted.put(member.memberName(), toSmithy(memberDocument, member)); + break; + } + } + yield Document.of(converted); + } + case LIST, SET -> { + var converted = new ArrayList(); + for (var item : document.asList()) { + converted.add(toSmithy(item, schema.listMember())); + } + yield Document.of(converted); + } + case MAP -> { + var converted = new HashMap(); + for (var entry : document.asStringMap().entrySet()) { + converted.put(entry.getKey(), toSmithy(entry.getValue(), schema.mapValueMember())); + } + yield Document.of(converted); + } + case DOCUMENT -> toSmithyOneOf(document, schema); + default -> document; + }; + } + + Document fromSmithy(Document document, Schema schema) { + if (document == null) { + return null; + } + if (!needsAdaptation(schema)) { + return document; + } + return switch (schema.type()) { + case BIG_DECIMAL -> Document.of(document.asBigDecimal().toString()); + case BIG_INTEGER -> Document.of(document.asBigInteger().toString()); + case BLOB -> Document.of(Base64.getEncoder().encodeToString(ByteBufferUtils.getBytes(document.asBlob()))); + case TIMESTAMP -> adaptTimestamp(document); + case STRUCTURE -> { + var converted = new HashMap(); + for (var member : schema.members()) { + var memberDocument = document.getMember(member.memberName()); + if (memberDocument != null) { + converted.put(member.memberName(), fromSmithy(memberDocument, member)); + } + } + yield Document.of(converted); + } + case UNION -> { + Document converted = Document.of(Map.of()); + for (var member : schema.members()) { + var memberDocument = document.getMember(member.memberName()); + if (memberDocument != null) { + converted = Document.of(Map.of( + member.memberName(), + fromSmithy(memberDocument, member))); + break; + } + } + yield converted; + } + case LIST, SET -> { + var converted = new ArrayList(); + for (var item : document.asList()) { + converted.add(fromSmithy(item, schema.listMember())); + } + yield Document.of(converted); + } + case MAP -> { + var converted = new HashMap(); + for (var entry : document.asStringMap().entrySet()) { + converted.put(entry.getKey(), fromSmithy(entry.getValue(), schema.mapValueMember())); + } + yield Document.of(converted); + } + case DOCUMENT -> fromSmithyOneOf(document, schema); + default -> document; + }; + } + + private Document toSmithyOneOf(Document document, Schema schema) { + var targetSchema = schema.isMember() ? schema.memberTarget() : schema; + var oneOf = targetSchema.getTrait(ONE_OF_TRAIT); + if (oneOf == null) { + return document; + } + + for (var member : oneOf.getMembers()) { + var memberDocument = document.getMember(member.getName()); + if (memberDocument != null) { + var converted = new HashMap(); + converted.put(oneOf.getDiscriminator(), Document.of(member.getTarget().toString())); + converted.putAll(toSmithy(memberDocument, schemaIndex.getSchema(member.getTarget())).asStringMap()); + return Document.of(converted); + } + } + return document; + } + + private Document fromSmithyOneOf(Document document, Schema schema) { + var targetSchema = schema.isMember() ? schema.memberTarget() : schema; + var oneOf = targetSchema.getTrait(ONE_OF_TRAIT); + if (oneOf == null) { + return document; + } + + var discriminator = document.getMember(oneOf.getDiscriminator()); + if (discriminator == null) { + return document; + } + + var shapeId = ShapeId.from(discriminator.asString()); + for (var member : oneOf.getMembers()) { + if (member.getTarget().equals(shapeId)) { + var converted = new HashMap<>( + fromSmithy(document, schemaIndex.getSchema(shapeId)).asStringMap()); + converted.remove(oneOf.getDiscriminator()); + return Document.of(Map.of(member.getName(), Document.of(converted))); + } + } + return document; + } + + private boolean needsAdaptation(Schema schema) { + return adaptationRequired.computeIfAbsent( + schema, + ignored -> needsAdaptation( + schema, + Collections.newSetFromMap(new IdentityHashMap<>()))); + } + + private boolean needsAdaptation(Schema schema, Set visiting) { + var target = schema.isMember() ? schema.memberTarget() : schema; + if (!visiting.add(schema)) { + return false; + } + try { + return switch (target.type()) { + case BIG_DECIMAL, BIG_INTEGER, BLOB, TIMESTAMP, DOCUMENT, UNION -> true; + case STRUCTURE -> target.members() + .stream() + .anyMatch(member -> needsAdaptation(member, visiting)); + case LIST, SET -> needsAdaptation(target.listMember(), visiting); + case MAP -> needsAdaptation(target.mapValueMember(), visiting); + default -> false; + }; + } finally { + visiting.remove(schema); + } + } + + private static Document badType(ShapeType from, ShapeType to) { + throw new IllegalArgumentException("Cannot convert from " + from + " to " + to); + } + + private static Document adaptTimestamp(Document document) { + if (document.isType(ShapeType.TIMESTAMP)) { + return Document.of(DATE_TIME.writeString(document.asTimestamp())); + } + if (document.isType(ShapeType.STRING)) { + var value = document.asString(); + try { + return Document.of(DATE_TIME.readFromString(value, false)); + } catch (RuntimeException e) { + return Document.of(HTTP_DATE.readFromString(value, false)); + } + } + return Document.of(EPOCH_SECONDS.readFromNumber(document.asNumber())); + } +} 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/StdioMcpClient.java similarity index 64% rename from mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioProxy.java rename to mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpClient.java index 94d7388d2f..cacb8b3964 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/StdioMcpClient.java @@ -16,34 +16,33 @@ import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; -import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import software.amazon.smithy.java.core.serde.document.Document; -import software.amazon.smithy.java.json.JsonCodec; 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.utils.SmithyUnstableApi; @SmithyUnstableApi -public final class StdioProxy extends McpServerProxy { - private static final InternalLogger LOG = InternalLogger.getLogger(StdioProxy.class); - private static final JsonCodec JSON_CODEC = JsonCodec.builder().build(); +public final class StdioMcpClient extends McpRemoteClient { + private static final InternalLogger LOG = InternalLogger.getLogger(StdioMcpClient.class); private final ProcessBuilder processBuilder; - private Process process; - private BufferedReader reader; - private BufferedWriter writer; + private volatile Process process; + private volatile BufferedReader reader; + private volatile BufferedWriter writer; private final Lock writeLock = new ReentrantLock(); private Thread responseReaderThread; private Thread errorReaderThread; - private final Map> pendingRequests = new ConcurrentHashMap<>(); + private final Map pendingRequests = new ConcurrentHashMap<>(); private volatile boolean running = false; private final String name; - private StdioProxy(Builder builder) { + private StdioMcpClient(Builder builder) { processBuilder = new ProcessBuilder(); processBuilder.command().add(builder.command); @@ -98,11 +97,11 @@ public Builder workingDirectory(File workingDirectory) { return this; } - public StdioProxy build() { + public StdioMcpClient build() { if (command == null || command.isEmpty()) { throw new IllegalArgumentException("Command must be provided"); } - return new StdioProxy(this); + return new StdioMcpClient(this); } } @@ -111,39 +110,36 @@ public static Builder builder() { } @Override - public CompletableFuture rpc(JsonRpcRequest request) { + protected JsonRpcResponse exchange(JsonRpcRequest request) { if (process == null || !process.isAlive()) { - CompletableFuture future = new CompletableFuture<>(); - future.completeExceptionally(new IllegalStateException("MCP server process is not running")); - return future; + throw new McpRemoteException("MCP server process is not running"); } // Notifications don't have an ID and don't expect a response if (request.getId() == null) { + String serializedRequest = McpJson.CODEC.serializeToString(request); try { writeLock.lock(); - String serializedRequest = JSON_CODEC.serializeToString(request); LOG.debug("Sending notification: {}", serializedRequest); writer.write(serializedRequest); writer.newLine(); writer.flush(); } catch (IOException e) { LOG.error("Error sending notification to MCP server", e); - return CompletableFuture.failedFuture( - new RuntimeException("Failed to send notification to MCP server: " + e.getMessage(), e)); + throw new McpRemoteException("Failed to send notification to MCP server", e); } finally { writeLock.unlock(); } - return CompletableFuture.completedFuture(null); + return null; } - String requestId = getStringRequestId(request.getId()); - CompletableFuture responseFuture = new CompletableFuture<>(); - pendingRequests.put(requestId, responseFuture); + String requestId = requestKey(request.getId()); + String serializedRequest = McpJson.CODEC.serializeToString(request); + var pending = new PendingResponse(); + pendingRequests.put(requestId, pending); try { writeLock.lock(); - String serializedRequest = JSON_CODEC.serializeToString(request); LOG.debug("Sending request ID {}: {}", requestId, serializedRequest); writer.write(serializedRequest); @@ -152,19 +148,18 @@ public CompletableFuture rpc(JsonRpcRequest request) { } catch (IOException e) { LOG.error("Error sending request to MCP server", e); pendingRequests.remove(requestId); - responseFuture.completeExceptionally( - new RuntimeException("Failed to send request to MCP server: " + e.getMessage(), e)); + throw new McpRemoteException("Failed to send request to MCP server", e); } finally { writeLock.unlock(); } - return responseFuture; + return pending.await(); } - private String getStringRequestId(Document id) { + static String requestKey(Document id) { return switch (id.type()) { - case STRING -> id.asString(); - case INTEGER -> Integer.toString(id.asInteger()); + case STRING -> "string:" + id.asString(); + case INTEGER, LONG, BIG_INTEGER -> "number:" + id.asBigInteger(); default -> throw new IllegalStateException("Unexpected value: " + id.type()); }; } @@ -207,19 +202,19 @@ public synchronized void start() { LOG.debug("Received response: {}", responseLine); var output = - JSON_CODEC.createDeserializer(responseLine.getBytes(StandardCharsets.UTF_8)) + McpJson.CODEC.createDeserializer(responseLine.getBytes(StandardCharsets.UTF_8)) .readDocument(); if (isNotification(output)) { notify(output.asShape(JsonRpcRequest.builder())); } else { JsonRpcResponse response = output.asShape(JsonRpcResponse.builder()); - String responseId = getStringRequestId(response.getId()); + String responseId = requestKey(response.getId()); LOG.debug("Processing response ID: {}", responseId); - CompletableFuture future = pendingRequests.remove(responseId); - if (future != null) { - future.complete(response); + PendingResponse pending = pendingRequests.remove(responseId); + if (pending != null) { + pending.complete(response); } else { notify(response); } @@ -236,8 +231,8 @@ public synchronized void start() { // Complete all pending requests with an exception if the reader exits if (!pendingRequests.isEmpty()) { - pendingRequests.forEach((id, future) -> future - .completeExceptionally(new RuntimeException("MCP server connection closed"))); + pendingRequests.forEach((id, pending) -> pending + .fail(new McpRemoteException("MCP server connection closed"))); pendingRequests.clear(); } }); @@ -248,51 +243,71 @@ public synchronized void start() { } @Override - public CompletableFuture shutdown() { - return CompletableFuture.runAsync(() -> { - running = false; - if (process != null && process.isAlive()) { - try { - // Complete all pending requests with exceptions - pendingRequests.forEach((id, future) -> future - .completeExceptionally(new RuntimeException("MCP server shutting down"))); - pendingRequests.clear(); - - // Close streams - if (writer != null) { - writer.close(); - } - if (reader != null) { - reader.close(); - } - - // Interrupt the response reader thread - if (responseReaderThread != null && responseReaderThread.isAlive()) { - responseReaderThread.interrupt(); - } - - if (errorReaderThread != null && errorReaderThread.isAlive()) { - errorReaderThread.interrupt(); - } + public void close() { + running = false; + if (process != null && process.isAlive()) { + try { + pendingRequests.forEach((id, pending) -> pending + .fail(new McpRemoteException("MCP server shutting down"))); + pendingRequests.clear(); - // Destroy the process - process.destroy(); + if (writer != null) { + writer.close(); + } + if (reader != null) { + reader.close(); + } + if (responseReaderThread != null && responseReaderThread.isAlive()) { + responseReaderThread.interrupt(); + } + if (errorReaderThread != null && errorReaderThread.isAlive()) { + errorReaderThread.interrupt(); + } - // Wait for termination with timeout - if (!process.waitFor(5, SECONDS)) { - // Force kill if it doesn't terminate gracefully - process.destroyForcibly(); - } - } catch (IOException | InterruptedException e) { - LOG.error("Error shutting down MCP server process", e); - Thread.currentThread().interrupt(); + process.destroy(); + if (!process.waitFor(5, SECONDS)) { + process.destroyForcibly(); } + } catch (IOException e) { + LOG.error("Error shutting down MCP server process", e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("Interrupted while shutting down MCP server", e); } - }); + } } @Override public String name() { return this.name; } + + private static final class PendingResponse { + private final BlockingQueue result = new ArrayBlockingQueue<>(1); + + void complete(JsonRpcResponse response) { + if (!result.offer(response)) { + throw new IllegalStateException("MCP request was already completed"); + } + } + + void fail(RuntimeException error) { + if (!result.offer(error)) { + throw new IllegalStateException("MCP request was already completed"); + } + } + + JsonRpcResponse await() { + try { + return switch (result.take()) { + case JsonRpcResponse response -> response; + case RuntimeException error -> throw error; + default -> throw new IllegalStateException("Unexpected pending MCP response"); + }; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("Interrupted while waiting for MCP response", e); + } + } + } } diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServer.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServer.java new file mode 100644 index 0000000000..d4ce88442b --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServer.java @@ -0,0 +1,183 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import software.amazon.smithy.java.core.schema.SerializableStruct; +import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.logging.InternalLogger; +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.server.Server; +import software.amazon.smithy.java.server.Service; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * MCP server using newline-delimited JSON-RPC over standard input and output. + * + *

Requests execute on virtual threads. Initialization is awaited before additional + * input is dispatched so protocol negotiation cannot race later requests. + */ +@SmithyUnstableApi +public final class StdioMcpServer implements Server { + private static final InternalLogger LOG = InternalLogger.getLogger(StdioMcpServer.class); + private static final byte[] TOOLS_CHANGED = """ + {"jsonrpc":"2.0","method":"notifications/tools/list_changed"} + """.getBytes(StandardCharsets.UTF_8); + + private final McpEngine engine; + private final Thread listener; + private final InputStream input; + private final OutputStream output; + private final McpSession session; + private final ExecutorService requests = Executors.newVirtualThreadPerTaskExecutor(); + private final CountDownLatch done = new CountDownLatch(1); + + StdioMcpServer(StdioMcpServerBuilder builder) { + engine = builder.engine; + session = engine.newSession(); + input = builder.input; + output = builder.output; + listener = Thread.ofPlatform() + .name("stdio-dispatcher") + .daemon() + .unstarted(() -> { + try { + listen(); + } catch (RuntimeException e) { + LOG.error("Error handling MCP input", e); + } finally { + done.countDown(); + } + }); + } + + private void listen() { + try (var reader = new BufferedReader(new InputStreamReader(input, StandardCharsets.UTF_8))) { + String line; + while ((line = reader.readLine()) != null) { + final JsonRpcRequest request; + try { + request = McpJson.CODEC.deserializeShape(line, JsonRpcRequest.builder()); + } catch (RuntimeException e) { + LOG.error("Error decoding MCP request", e); + write(parseError()); + continue; + } + + var task = requests.submit(() -> handleRequest(request)); + if (McpMethod.Standard.INITIALIZE.wireName().equals(request.getMethod())) { + try { + task.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } catch (ExecutionException e) { + LOG.error("Error dispatching MCP initialize request", e.getCause()); + } + } + } + } catch (IOException e) { + LOG.error("Error reading MCP input", e); + } finally { + requests.shutdown(); + } + } + + private void handleRequest(JsonRpcRequest request) { + var outcome = engine.execute(request, session, null, McpTransportContext.STDIO); + var response = engine.encode(outcome); + if (response != null) { + write(response); + } + } + + public void refreshTools() { + try { + synchronized (output) { + output.write(TOOLS_CHANGED); + output.flush(); + } + } catch (IOException e) { + LOG.error("Failed to write tools-changed notification", e); + } + } + + public void addService(String id, Service service) { + engine.addService(id, service); + refreshTools(); + } + + public void addRemoteClient(McpRemoteClient client) { + engine.addRemoteClient(client); + refreshTools(); + } + + public boolean containsServer(String id) { + return engine.containsServer(id); + } + + private void write(SerializableStruct shape) { + var bytes = McpJson.CODEC.serialize(shape); + synchronized (output) { + try { + if (bytes.hasArray()) { + output.write(bytes.array(), bytes.arrayOffset() + bytes.position(), bytes.remaining()); + } else { + output.write(ByteBufferUtils.getBytes(bytes)); + } + output.write('\n'); + output.flush(); + } catch (IOException e) { + LOG.error("Error writing MCP output", e); + } + } + } + + private JsonRpcResponse parseError() { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .error(JsonRpcErrorResponse.builder() + .code(-32700) + .message("Parse error") + .build()) + .build(); + } + + @Override + public void start() { + engine.bindTransport(this::write, this::write); + listener.start(); + } + + @Override + public CompletableFuture shutdown() { + requests.shutdownNow(); + engine.close(); + return CompletableFuture.completedFuture(null); + } + + public void awaitCompletion() throws InterruptedException { + done.await(); + requests.awaitTermination(30, TimeUnit.SECONDS); + } + + public static StdioMcpServerBuilder builder() { + return new StdioMcpServerBuilder(); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServerBuilder.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServerBuilder.java new file mode 100644 index 0000000000..f553e19070 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServerBuilder.java @@ -0,0 +1,159 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.io.InputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import software.amazon.smithy.java.server.Service; +import software.amazon.smithy.utils.SmithyUnstableApi; + +@SmithyUnstableApi +public final class StdioMcpServerBuilder { + InputStream input; + OutputStream output; + McpEngine engine; + + private final Map services = new HashMap<>(); + private final List remoteClients = new ArrayList<>(); + private final Map protocols = new LinkedHashMap<>(); + private final Map protocolOverrides = new LinkedHashMap<>(); + private McpInterceptor interceptor = McpInterceptor.NOOP; + private String name = "mcp-server"; + private String version = "1.0.0"; + private ToolFilter toolFilter = (server, tool) -> true; + private McpMetricsObserver metricsObserver; + private boolean discoverProtocols = true; + + StdioMcpServerBuilder() {} + + public StdioMcpServerBuilder stdio() { + input = System.in; + output = System.out; + return this; + } + + public StdioMcpServerBuilder input(InputStream input) { + this.input = input; + return this; + } + + public StdioMcpServerBuilder output(OutputStream output) { + this.output = output; + return this; + } + + /** + * Uses a prebuilt engine instead of constructing one from this builder's source options. + */ + public StdioMcpServerBuilder engine(McpEngine engine) { + this.engine = Objects.requireNonNull(engine, "engine"); + return this; + } + + public StdioMcpServerBuilder name(String name) { + this.name = Objects.requireNonNull(name, "name"); + return this; + } + + public StdioMcpServerBuilder version(String version) { + this.version = Objects.requireNonNull(version, "version"); + return this; + } + + public StdioMcpServerBuilder addService(String id, Service service) { + services.put(id, service); + return this; + } + + public StdioMcpServerBuilder addServices(Map services) { + this.services.putAll(services); + return this; + } + + public StdioMcpServerBuilder addRemoteClient(McpRemoteClient... clients) { + remoteClients.addAll(Arrays.asList(clients)); + return this; + } + + public StdioMcpServerBuilder toolFilter(ToolFilter filter) { + toolFilter = Objects.requireNonNull(filter, "filter"); + return this; + } + + public StdioMcpServerBuilder metricsObserver(McpMetricsObserver observer) { + metricsObserver = observer; + return this; + } + + public StdioMcpServerBuilder interceptor(McpInterceptor interceptor) { + this.interceptor = Objects.requireNonNull(interceptor, "interceptor"); + return this; + } + + public StdioMcpServerBuilder addProtocol(ExtensionMcpProtocol protocol) { + putProtocol(protocols, protocol, "protocol"); + return this; + } + + public StdioMcpServerBuilder overrideProtocol(ExtensionMcpProtocol protocol) { + putProtocol(protocolOverrides, protocol, "protocol override"); + return this; + } + + public StdioMcpServerBuilder discoverProtocols(boolean discoverProtocols) { + this.discoverProtocols = discoverProtocols; + return this; + } + + public StdioMcpServer build() { + Objects.requireNonNull(input, "MCP server input stream is required"); + Objects.requireNonNull(output, "MCP server output stream is required"); + if (engine == null) { + if (services.isEmpty() && remoteClients.isEmpty()) { + throw new IllegalArgumentException("MCP server requires an engine, service, or remote client"); + } + + var engineBuilder = McpEngine.builder() + .services(services) + .remoteClients(remoteClients) + .name(name) + .version(version) + .toolFilter(toolFilter) + .metricsObserver(metricsObserver) + .interceptor(interceptor) + .discoverProtocols(discoverProtocols); + protocols.values().forEach(engineBuilder::addProtocol); + protocolOverrides.values().forEach(engineBuilder::overrideProtocol); + engine = engineBuilder.build(); + } else if (!services.isEmpty() + || !remoteClients.isEmpty() + || !protocols.isEmpty() + || !protocolOverrides.isEmpty()) { + throw new IllegalStateException("Cannot combine a prebuilt engine with builder-managed sources"); + } + return new StdioMcpServer(this); + } + + private void putProtocol( + Map destination, + ExtensionMcpProtocol protocol, + String kind + ) { + Objects.requireNonNull(protocol, kind); + Objects.requireNonNull(protocol.id(), kind + " id"); + if (destination.put(protocol.id(), protocol) != null) { + throw new IllegalArgumentException( + "Duplicate MCP " + kind + ": " + protocol.id().identifier()); + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/UnknownProtocolVersion.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/UnknownProtocolVersion.java new file mode 100644 index 0000000000..c6dd10f944 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/UnknownProtocolVersion.java @@ -0,0 +1,22 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server; + +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * A protocol version not built into the core library. + * + *

The engine may still support this version through a registered + * {@link ExtensionMcpProtocol}. + */ +@SmithyUnstableApi +public record UnknownProtocolVersion(String identifier) implements ProtocolVersion { + public UnknownProtocolVersion { + Objects.requireNonNull(identifier, "identifier"); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/package-info.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/package-info.java index 0221e2c326..d9e83b7eb5 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/package-info.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/package-info.java @@ -1,5 +1,16 @@ /** - * MCP server implementation for exposing Smithy services as tools. + * Extensible MCP execution, protocol, client, and transport support for exposing + * Smithy services as tools. + * + *

{@link software.amazon.smithy.java.mcp.server.McpEngine} is the blocking, + * transport-independent core. Standard calls and outcomes use sealed typed + * hierarchies. Custom methods are added with + * {@link software.amazon.smithy.java.mcp.server.McpExtensionMethod}, and external + * protocol versions implement + * {@link software.amazon.smithy.java.mcp.server.ExtensionMcpProtocol} directly or + * through {@link software.amazon.smithy.java.mcp.server.McpProtocolProvider}. + * Stdio and HTTP adapters own transport concerns and concurrency. + * *

This package is under development and is not intended for use in production. */ @SmithyUnstableApi 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/HttpMcpClientTest.java similarity index 78% rename from mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java rename to mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpClientTest.java index bf7542ce61..46614e4c48 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/HttpMcpClientTest.java @@ -15,8 +15,9 @@ import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -35,11 +36,11 @@ import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.model.shapes.ShapeType; -class HttpMcpProxyTest { +class HttpMcpClientTest { private static final JsonCodec JSON_CODEC = JsonCodec.builder().build(); private HttpServer mockServer; - private HttpMcpProxy proxy; + private HttpMcpClient proxy; private String serverUrl; @BeforeEach @@ -51,7 +52,7 @@ void setUp() throws IOException { mockServer.createContext("/mcp", new MockMcpHandler()); mockServer.start(); - proxy = HttpMcpProxy.builder() + proxy = HttpMcpClient.builder() .endpoint(serverUrl) .name("Test MCP") .build(); @@ -63,21 +64,21 @@ void tearDown() { mockServer.stop(0); } if (proxy != null) { - proxy.shutdown().join(); + proxy.close(); } } @Test void testBuilderValidation() { - assertThrows(IllegalArgumentException.class, () -> HttpMcpProxy.builder().build()); + assertThrows(IllegalArgumentException.class, () -> HttpMcpClient.builder().build()); - assertThrows(IllegalArgumentException.class, () -> HttpMcpProxy.builder().endpoint("").build()); + assertThrows(IllegalArgumentException.class, () -> HttpMcpClient.builder().endpoint("").build()); } @Test void testBuilderRejectsSignerAndAuthSchemeTogether() { assertThrows(IllegalArgumentException.class, - () -> HttpMcpProxy.builder() + () -> HttpMcpClient.builder() .endpoint(serverUrl) .signer((request, identity, context) -> new SignResult<>(request)) .authScheme(new TestAuthScheme()) @@ -88,7 +89,7 @@ void testBuilderRejectsSignerAndAuthSchemeTogether() { @Test void testBuilderRejectsAuthSchemeWithoutIdentityResolver() { assertThrows(IllegalArgumentException.class, - () -> HttpMcpProxy.builder() + () -> HttpMcpClient.builder() .endpoint(serverUrl) .authScheme(new TestAuthScheme()) .build()); @@ -97,7 +98,7 @@ void testBuilderRejectsAuthSchemeWithoutIdentityResolver() { @Test void testBuilderRejectsIdentityResolverWithoutAuthScheme() { assertThrows(IllegalArgumentException.class, - () -> HttpMcpProxy.builder() + () -> HttpMcpClient.builder() .endpoint(serverUrl) .identityResolver(TestIdentityResolver.INSTANCE) .build()); @@ -119,7 +120,7 @@ void testAuthSchemeSignsRequest() throws IOException { exchange.close(); }); - HttpMcpProxy authProxy = HttpMcpProxy.builder() + HttpMcpClient authProxy = HttpMcpClient.builder() .endpoint(serverUrl) .authScheme(new TestAuthScheme()) .identityResolver(TestIdentityResolver.INSTANCE) @@ -131,12 +132,12 @@ void testAuthSchemeSignsRequest() throws IOException { .jsonrpc("2.0") .build(); - JsonRpcResponse response = authProxy.rpc(request).join(); + JsonRpcResponse response = authProxy.exchange(request); assertNotNull(response); assertEquals("signed", response.getResult().asString()); assertEquals("test-token", capturedHeader[0]); - authProxy.shutdown().join(); + authProxy.close(); } @Test @@ -158,7 +159,7 @@ void testAuthSchemeReceivesSignerContext() throws IOException { Context signerCtx = Context.create(); signerCtx.put(TestAuthScheme.REGION_KEY, "us-west-2"); - HttpMcpProxy authProxy = HttpMcpProxy.builder() + HttpMcpClient authProxy = HttpMcpClient.builder() .endpoint(serverUrl) .authScheme(new TestAuthScheme()) .identityResolver(TestIdentityResolver.INSTANCE) @@ -171,28 +172,28 @@ void testAuthSchemeReceivesSignerContext() throws IOException { .jsonrpc("2.0") .build(); - JsonRpcResponse response = authProxy.rpc(request).join(); + JsonRpcResponse response = authProxy.exchange(request); assertNotNull(response); assertEquals("us-west-2", capturedRegion[0]); - authProxy.shutdown().join(); + authProxy.close(); } @Test void testBuilderWithCustomName() { - HttpMcpProxy customProxy = HttpMcpProxy.builder() + HttpMcpClient customProxy = HttpMcpClient.builder() .endpoint(serverUrl) .name("Custom Name") .build(); assertEquals("Custom Name", customProxy.name()); - customProxy.shutdown().join(); + customProxy.close(); } @Test void testBuilderWithHeaders() { Map headers = Map.of("Authorization", "Bearer token"); - HttpMcpProxy proxyWithHeaders = HttpMcpProxy.builder() + HttpMcpClient proxyWithHeaders = HttpMcpClient.builder() .endpoint(serverUrl) .signer((request, identity, context) -> { var r = request.toModifiable(); @@ -204,13 +205,13 @@ void testBuilderWithHeaders() { .build(); assertNotNull(proxyWithHeaders); - proxyWithHeaders.shutdown().join(); + proxyWithHeaders.close(); } @Test void testBuilderWithDynamicHeaders() { int[] counter = {0}; - HttpMcpProxy proxyWithDynamicHeaders = HttpMcpProxy.builder() + HttpMcpClient proxyWithDynamicHeaders = HttpMcpClient.builder() .endpoint(serverUrl) .signer((request, identity, context) -> { var r = request.toModifiable(); @@ -222,27 +223,27 @@ void testBuilderWithDynamicHeaders() { .build(); assertNotNull(proxyWithDynamicHeaders); - proxyWithDynamicHeaders.shutdown().join(); + proxyWithDynamicHeaders.close(); } @Test void testDefaultName() { - HttpMcpProxy defaultProxy = HttpMcpProxy.builder() + HttpMcpClient defaultProxy = HttpMcpClient.builder() .endpoint(serverUrl) .build(); assertEquals("localhost", defaultProxy.name()); - defaultProxy.shutdown().join(); + defaultProxy.close(); } @Test void testSanitizedName() { - HttpMcpProxy sanitizedProxy = HttpMcpProxy.builder() + HttpMcpClient sanitizedProxy = HttpMcpClient.builder() .endpoint("http://api.example.com:8080/path") .build(); assertEquals("api-example-com", sanitizedProxy.name()); - sanitizedProxy.shutdown().join(); + sanitizedProxy.close(); } @Test @@ -253,8 +254,7 @@ void testRpcCall() { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); assertNotNull(response); assertEquals("2.0", response.getJsonrpc()); @@ -264,10 +264,24 @@ void testRpcCall() { @Test void testRpcWithNullRequest() { - CompletableFuture future = proxy.rpc(null); + assertThrows(McpRemoteException.class, () -> proxy.exchange(null)); + } + + @Test + void testNotificationAcceptsEmptySuccessfulResponse() throws IOException { + mockServer.removeContext("/mcp"); + mockServer.createContext("/mcp", exchange -> { + exchange.getRequestBody().readAllBytes(); + exchange.sendResponseHeaders(202, -1); + exchange.close(); + }); - ExecutionException exception = assertThrows(ExecutionException.class, future::get); - assertTrue(exception.getCause() instanceof NullPointerException); + var notification = JsonRpcRequest.builder() + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .jsonrpc("2.0") + .build(); + + assertNull(proxy.exchange(notification)); } @Test @@ -284,12 +298,12 @@ void testRpcHttpError() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); assertNotNull(response); assertNotNull(response.getError()); - assertEquals(500, response.getError().getCode()); + assertEquals(1, response.getId().asInteger()); + assertEquals(-32000, response.getError().getCode()); assertTrue(response.getError().getMessage().contains("HTTP 500")); } @@ -297,7 +311,7 @@ void testRpcHttpError() throws IOException { void testStartAndShutdown() { assertDoesNotThrow(() -> { proxy.start(); - proxy.shutdown().join(); + proxy.close(); }); } @@ -313,8 +327,7 @@ void testSseStreamingResponse() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); assertNotNull(response); assertEquals("2.0", response.getJsonrpc()); @@ -342,7 +355,7 @@ void testSseStreamingWithNotifications() throws IOException { notification -> {}, // Old-style consumer (not used) notification -> capturedNotification[0] = notification, // Request notification consumer initRequest, - ProtocolVersion.defaultVersion()); + BuiltInProtocols.protocol(ProtocolVersion.defaultVersion())); JsonRpcRequest request = JsonRpcRequest.builder() .method("test/streaming") @@ -350,8 +363,7 @@ void testSseStreamingWithNotifications() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); // Verify final response assertNotNull(response); @@ -364,6 +376,37 @@ void testSseStreamingWithNotifications() throws IOException { assertNull(capturedNotification[0].getId()); } + @Test + void testSseNotificationsAreDeliveredBeforeTheFinalResponse() throws IOException { + var notificationObserved = new CountDownLatch(1); + var observedBeforeFinal = new AtomicBoolean(); + + mockServer.removeContext("/mcp"); + mockServer.createContext( + "/mcp", + new LiveSseNotificationHandler(notificationObserved, observedBeforeFinal)); + + var initRequest = JsonRpcRequest.builder() + .method(McpMethod.Standard.INITIALIZE.wireName()) + .id(Document.of(0)) + .jsonrpc("2.0") + .build(); + proxy.initialize( + ignored -> {}, + ignored -> notificationObserved.countDown(), + initRequest, + BuiltInProtocols.protocol(ProtocolVersion.defaultVersion())); + + var response = proxy.exchange(JsonRpcRequest.builder() + .method("test/streaming") + .id(Document.of(1)) + .jsonrpc("2.0") + .build()); + + assertEquals("final result", response.getResult().asString()); + assertTrue(observedBeforeFinal.get()); + } + @Test void testSseStreamingWithoutFinalResponse() throws IOException { // Set up SSE handler that doesn't send a final response @@ -376,8 +419,7 @@ void testSseStreamingWithoutFinalResponse() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); // Should return an error response assertNotNull(response); @@ -398,8 +440,7 @@ void testSseStreamingMalformedJson() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); // Should return an error response assertNotNull(response); @@ -420,8 +461,7 @@ void testSseStreamingWithMethodInToolResponse() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future = proxy.rpc(request); - JsonRpcResponse response = future.join(); + JsonRpcResponse response = proxy.exchange(request); // Should correctly parse as a response, not a notification assertNotNull(response); @@ -461,8 +501,7 @@ void testSessionIdHandling() throws IOException { Document.of("1.0.0")))))) .build(); - CompletableFuture future1 = proxy.rpc(request1); - JsonRpcResponse response1 = future1.join(); + JsonRpcResponse response1 = proxy.exchange(request1); assertNotNull(response1); assertEquals("session-created", response1.getResult().asString()); @@ -474,8 +513,7 @@ void testSessionIdHandling() throws IOException { .jsonrpc("2.0") .build(); - CompletableFuture future2 = proxy.rpc(request2); - JsonRpcResponse response2 = future2.join(); + JsonRpcResponse response2 = proxy.exchange(request2); assertNotNull(response2); assertEquals("session-valid", response2.getResult().asString()); @@ -521,6 +559,78 @@ public void handle(HttpExchange exchange) throws IOException { } } + private static class LiveSseNotificationHandler implements HttpHandler { + private final CountDownLatch notificationObserved; + private final AtomicBoolean observedBeforeFinal; + + LiveSseNotificationHandler( + CountDownLatch notificationObserved, + AtomicBoolean observedBeforeFinal + ) { + this.notificationObserved = notificationObserved; + this.observedBeforeFinal = observedBeforeFinal; + } + + @Override + public void handle(HttpExchange exchange) throws IOException { + var requestBytes = exchange.getRequestBody().readAllBytes(); + var request = JsonRpcRequest.builder() + .deserialize(JSON_CODEC.createDeserializer(requestBytes)) + .build(); + + if (McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName().equals(request.getMethod())) { + exchange.sendResponseHeaders(202, -1); + exchange.close(); + return; + } + if (McpMethod.Standard.INITIALIZE.wireName().equals(request.getMethod())) { + writeJsonResponse(exchange, request.getId(), Document.of(Map.of())); + return; + } + + exchange.getResponseHeaders().set("Content-Type", "text/event-stream; charset=utf-8"); + exchange.sendResponseHeaders(200, 0); + try (var output = exchange.getResponseBody()) { + output.write(("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\"," + + "\"params\":{\"progress\":50}}\n\n") + .getBytes(StandardCharsets.UTF_8)); + output.flush(); + + try { + observedBeforeFinal.set(notificationObserved.await(2, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + + output.write( + "data: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":\"final result\"}\n\n" + .getBytes(StandardCharsets.UTF_8)); + } finally { + exchange.close(); + } + } + + private void writeJsonResponse( + HttpExchange exchange, + Document id, + Document result + ) throws IOException { + var response = JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(id) + .result(result) + .build(); + var body = JSON_CODEC.serializeToString(response).getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().set("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (var output = exchange.getResponseBody()) { + output.write(body); + } finally { + exchange.close(); + } + } + } + private static class SseStreamingNoFinalResponseHandler 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/McpArchitectureTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpArchitectureTest.java new file mode 100644 index 0000000000..5d57ae86f4 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpArchitectureTest.java @@ -0,0 +1,459 @@ +/* + * 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.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.net.URLClassLoader; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; + +class McpArchitectureTest { + + @Test + void everyKnownVersionSelectsItsOwnProtocol() { + for (var version : KnownProtocolVersion.values()) { + assertEquals(version, BuiltInProtocols.protocol(version).version()); + } + } + + @Test + void unsupportedProtocolOperationUsesDefaultBehavior() { + try (var engine = McpEngine.builder().build()) { + var outcome = engine.execute( + new McpCall.ReadResource( + Document.of(1), + "test://resource", + McpMetadata.EMPTY), + context(KnownProtocolVersion.V2025_11_25)); + + var failure = assertInstanceOf(McpOutcome.Failure.class, outcome); + assertEquals(-32601, failure.error().code()); + assertEquals("Method not found: resources/read", failure.error().message()); + } + } + + @Test + void userUnsupportedOperationExceptionIsAnInternalError() { + var interceptor = new McpInterceptor() { + @Override + public void readBeforeExecution(McpExecutionContext context) { + throw new UnsupportedOperationException("user code failed"); + } + }; + + try (var engine = McpEngine.builder().interceptor(interceptor).build()) { + var outcome = engine.execute( + new McpCall.Ping(Document.of(1), McpMetadata.EMPTY), + context(KnownProtocolVersion.V2025_11_25)); + + var failure = assertInstanceOf(McpOutcome.Failure.class, outcome); + assertEquals(-32603, failure.error().code()); + assertEquals("Internal error", failure.error().message()); + } + } + + @Test + void afterExecutionFailureDoesNotReplaceTheOriginalFailure() { + var original = new IllegalStateException("original"); + var after = new IllegalArgumentException("after"); + var observed = new AtomicReference(); + var interceptor = new McpInterceptor() { + @Override + public void readBeforeExecution(McpExecutionContext context) { + throw original; + } + + @Override + public void readAfterExecution( + McpExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + throw after; + } + + @Override + public McpOutcome modifyAfterExecution( + McpExecutionContext context, + McpOutcome outcome, + RuntimeException error + ) { + observed.set(error); + throw error; + } + }; + + try (var engine = McpEngine.builder().interceptor(interceptor).build()) { + engine.execute( + new McpCall.Ping(Document.of(1), McpMetadata.EMPTY), + context(KnownProtocolVersion.V2025_11_25)); + } + + assertSame(original, observed.get()); + assertSame(after, original.getSuppressed()[0]); + } + + @Test + void statelessProtocolDoesNotAccidentallyInheritLegacyPing() { + try (var engine = McpEngine.builder().build()) { + var metadata = new McpMetadata( + KnownProtocolVersion.V2026_07_28, + null, + Document.of(Map.of()), + Map.of()); + var outcome = engine.execute( + new McpCall.Ping(Document.of(1), metadata), + context(KnownProtocolVersion.V2026_07_28)); + + var failure = assertInstanceOf(McpOutcome.Failure.class, outcome); + assertEquals(-32601, failure.error().code()); + } + } + + @Test + void statelessProtocolRejectsRemovedMethodBeforeCapabilityParameterValidation() { + try (var engine = McpEngine.builder().build()) { + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.LOGGING_SET_LEVEL.wireName()) + .params(statelessParams()) + .build(); + + var response = engine.execute(request, KnownProtocolVersion.V2026_07_28); + + assertEquals(-32601, response.getError().getCode()); + } + } + + @Test + void unsupportedVersionReportsRequestedAndSupportedVersions() { + try (var engine = McpEngine.builder().build()) { + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.SERVER_DISCOVER.wireName()) + .build(); + + var response = engine.execute(request, new UnknownProtocolVersion("v999.0.0")); + + assertEquals(-32022, response.getError().getCode()); + assertEquals("v999.0.0", response.getError().getData().getMember("requested").asString()); + assertEquals( + KnownProtocolVersion.supportedIdentifiers().size(), + response.getError().getData().getMember("supported").asList().size()); + } + } + + @Test + void extensionMethodsDecodeExecuteAndEncodeTypedParameters() { + var extension = new EchoExtension(); + try (var engine = McpEngine.builder().addExtension(extension).build()) { + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of("extension-id")) + .method(extension.method()) + .params(Document.of(Map.of("value", Document.of("hello")))) + .build(); + + var response = engine.execute(request, KnownProtocolVersion.V2025_11_25); + assertNull(response.getError()); + assertEquals("hello", response.getResult().getMember("echo").asString()); + + var call = new McpCall.ExtensionCall<>( + Document.of(2), + extension, + new EchoParameters("outbound"), + McpMetadata.EMPTY); + var encoded = new McpWireCodec(Map.of(extension.method(), extension)).encode(call); + assertEquals(extension.method(), encoded.getMethod()); + assertEquals("outbound", encoded.getParams().getMember("value").asString()); + } + } + + @Test + void extensionCannotReplaceAStandardMethod() { + var extension = new EchoExtension() { + @Override + public String method() { + return McpMethod.Standard.PING.wireName(); + } + }; + assertThrows(IllegalArgumentException.class, () -> McpEngine.builder().addExtension(extension)); + } + + @Test + void extensionProtocolNegotiatesAndDispatchesThroughTheRegistry() { + var protocol = protocol("2099-01-01", Set.of(McpMethod.Standard.PING)); + try (var engine = McpEngine.builder() + .discoverProtocols(false) + .addProtocol(protocol) + .build()) { + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.PING.wireName()) + .build(); + + var response = engine.execute(request, new UnknownProtocolVersion(protocol.id().identifier())); + + assertNull(response.getError()); + assertEquals(Map.of(), response.getResult().asStringMap()); + } + } + + @Test + void initializeNegotiatesAnExtensionProtocol() { + var protocol = protocol("2099-01-01", Set.of(McpMethod.Standard.INITIALIZE)); + try (var engine = McpEngine.builder() + .discoverProtocols(false) + .addProtocol(protocol) + .build()) { + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of( + "protocolVersion", + Document.of(protocol.id().identifier()), + "capabilities", + Document.of(Map.of()), + "clientInfo", + Document.of(Map.of())))) + .build(); + + var response = engine.execute(request, null); + + assertNull(response.getError()); + assertEquals( + protocol.id().identifier(), + response.getResult().getMember("protocolVersion").asString()); + } + } + + @Test + void initializeFallsBackToLatestInitializationCapableProtocol() { + try (var engine = McpEngine.builder().build()) { + for (var requested : List.of( + new UnknownProtocolVersion("2099-01-01"), + KnownProtocolVersion.V2026_07_28)) { + var response = engine.execute(initializeRequest(requested), null); + + assertNull(response.getError()); + assertEquals( + KnownProtocolVersion.V2025_11_25.identifier(), + response.getResult().getMember("protocolVersion").asString()); + } + } + } + + @Test + void protocolErrorsDoNotProduceResponsesForNotifications() { + try (var engine = McpEngine.builder().build()) { + var notification = JsonRpcRequest.builder() + .jsonrpc("2.0") + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .build(); + + var response = engine.execute( + notification, + new UnknownProtocolVersion("unsupported-version")); + + assertNull(response); + } + } + + @Test + void programmaticProtocolConflictFailsWithoutAnOverride() { + var protocol = protocol( + KnownProtocolVersion.V2025_11_25.identifier(), + Set.of(McpMethod.Standard.PING)); + + var builder = McpEngine.builder() + .discoverProtocols(false) + .addProtocol(protocol); + + var error = assertThrows(IllegalStateException.class, builder::build); + assertTrue(error.getMessage().contains(protocol.id().identifier())); + } + + @Test + void explicitOverrideReplacesABuiltInProtocol() { + var protocol = protocol(KnownProtocolVersion.V2025_11_25.identifier(), Set.of()); + try (var engine = McpEngine.builder() + .discoverProtocols(false) + .overrideProtocol(protocol) + .build()) { + var response = engine.execute( + JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.PING.wireName()) + .build(), + KnownProtocolVersion.V2025_11_25); + + assertEquals(-32601, response.getError().getCode()); + } + } + + @Test + void spiConflictWithANewBuiltInFailsUnlessExplicitlyOverridden() { + var protocol = protocol( + KnownProtocolVersion.V2025_11_25.identifier(), + Set.of(McpMethod.Standard.PING)); + McpProtocolProvider provider = () -> List.of(protocol); + + var error = assertThrows( + IllegalStateException.class, + () -> McpProtocolRegistry.create(List.of(), List.of(), List.of(provider))); + assertTrue(error.getMessage().contains("SPI provider")); + + var registry = McpProtocolRegistry.create( + List.of(), + List.of(protocol), + List.of(provider)); + assertSame(protocol, registry.require(KnownProtocolVersion.V2025_11_25)); + } + + @Test + void conflictingSpiProvidersFailDeterministically() { + var first = protocol("2099-01-01", Set.of(McpMethod.Standard.PING)); + var second = protocol("2099-01-01", Set.of(McpMethod.Standard.TOOLS_LIST)); + McpProtocolProvider firstProvider = () -> List.of(first); + McpProtocolProvider secondProvider = () -> List.of(second); + + var error = assertThrows( + IllegalStateException.class, + () -> McpProtocolRegistry.create( + List.of(), + List.of(), + List.of(firstProvider, secondProvider))); + + assertTrue(error.getMessage().contains("2099-01-01")); + } + + @Test + void discoversProtocolProvidersWithServiceLoader(@TempDir Path temporaryDirectory) throws Exception { + var serviceDirectory = temporaryDirectory.resolve("META-INF/services"); + Files.createDirectories(serviceDirectory); + Files.writeString( + serviceDirectory.resolve(McpProtocolProvider.class.getName()), + TestProtocolProvider.class.getName()); + + try (var classLoader = new URLClassLoader( + new java.net.URL[] {temporaryDirectory.toUri().toURL()}, + getClass().getClassLoader())) { + var registry = McpProtocolRegistry.create(List.of(), List.of(), classLoader); + + assertEquals( + TestProtocolProvider.ID, + registry.require(new UnknownProtocolVersion(TestProtocolProvider.ID)).id().identifier()); + } + } + + @Test + void protocolRegistryReturnsSingletonImplementations() { + assertSame( + BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28), + BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28)); + } + + private McpRequestContext context(KnownProtocolVersion version) { + return new McpRequestContext(version, McpTransportContext.STDIO, Context.create()); + } + + private TestProtocol protocol(String identifier, Set methods) { + return new TestProtocol(McpProtocolId.of(identifier), methods); + } + + private Document statelessParams() { + return Document.of(Map.of( + "_meta", + Document.of(Map.of( + McpWireNames.PROTOCOL_VERSION, + Document.of(KnownProtocolVersion.V2026_07_28.identifier()), + McpWireNames.CLIENT_CAPABILITIES, + Document.of(Map.of()))))); + } + + private JsonRpcRequest initializeRequest(ProtocolVersion requestedVersion) { + return JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of( + "protocolVersion", + Document.of(requestedVersion.identifier()), + "capabilities", + Document.of(Map.of()), + "clientInfo", + Document.of(Map.of())))) + .build(); + } + + private record EchoParameters(String value) {} + + private record TestProtocol( + McpProtocolId id, + Set supportedMethods) implements ExtensionMcpProtocol { + private TestProtocol { + supportedMethods = Set.copyOf(supportedMethods); + } + } + + public static final class TestProtocolProvider implements McpProtocolProvider { + private static final String ID = "2099-service-loader"; + + @Override + public List protocols() { + return List.of(new TestProtocol( + McpProtocolId.of(ID), + Set.of(McpMethod.Standard.PING))); + } + } + + private static class EchoExtension implements McpExtensionMethod { + @Override + public String method() { + return "example/echo"; + } + + @Override + public EchoParameters decode(Document params) { + return new EchoParameters(params.getMember("value").asString()); + } + + @Override + public Document encode(EchoParameters params) { + return Document.of(Map.of("value", Document.of(params.value()))); + } + + @Override + public McpOutcome execute( + McpCall.ExtensionCall call, + McpRequestContext context + ) { + return new McpOutcome.Success( + call.id(), + Document.of(Map.of("echo", Document.of(call.parameters().value())))); + } + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpCatalogTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpCatalogTest.java new file mode 100644 index 0000000000..3bfb0d0c88 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpCatalogTest.java @@ -0,0 +1,319 @@ +/* + * 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.IntStream; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonObjectSchema; +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.PromptInfo; +import software.amazon.smithy.java.mcp.model.ToolInfo; + +class McpCatalogTest { + + @Test + void failingRemoteDoesNotAbortInitializationOrHideHealthyRemotes() { + var failing = new TestRemoteClient("failing") { + @Override + public List listTools() { + throw new McpRemoteException("unavailable"); + } + }; + var healthy = new TestRemoteClient("healthy") { + @Override + public List listTools() { + return List.of(tool("healthy-tool")); + } + }; + + try (var engine = McpEngine.builder() + .remoteClients(List.of(failing, healthy)) + .build()) { + var initialize = engine.execute(initializeRequest(), KnownProtocolVersion.V2025_11_25); + assertNull(initialize.getError()); + + var response = engine.execute( + request(2, McpMethod.Standard.TOOLS_LIST.wireName()), + KnownProtocolVersion.V2025_11_25); + var tools = response.getResult().asShape(ListToolsResult.builder()).getTools(); + + assertEquals(List.of("healthy-tool"), tools.stream().map(ToolInfo::getName).toList()); + } + } + + @Test + void remoteCatalogRefreshesRunInParallel() { + var entered = new CountDownLatch(2); + var firstTimedOut = new AtomicBoolean(); + var secondTimedOut = new AtomicBoolean(); + var first = blockingToolClient("first", "first-tool", entered, firstTimedOut); + var second = blockingToolClient("second", "second-tool", entered, secondTimedOut); + + try (var catalog = new McpCatalog(Map.of(), List.of(first, second))) { + catalog.ensureRemoteCatalogLoaded(); + + assertFalse(firstTimedOut.get()); + assertFalse(secondTimedOut.get()); + assertEquals(2, catalog.snapshot().tools().size()); + } + } + + @Test + void remoteIoDoesNotBlockCatalogReads() throws Exception { + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var remote = new TestRemoteClient("blocking") { + @Override + public List listTools() { + entered.countDown(); + try { + assertTrue(release.await(5, SECONDS)); + return List.of(tool("blocking-tool")); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("interrupted", e); + } + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + var refresh = Thread.ofVirtual().start(catalog::ensureRemoteCatalogLoaded); + try { + assertTrue(entered.await(2, SECONDS)); + assertTimeoutPreemptively( + Duration.ofMillis(500), + () -> assertTrue(catalog.containsServer("blocking"))); + } finally { + release.countDown(); + refresh.join(); + } + } + } + + @Test + void concurrentRemoteAdditionsPublishWithoutLostUpdates() throws Exception { + var clientCount = 32; + var ready = new CountDownLatch(clientCount); + var start = new CountDownLatch(1); + + try (var catalog = new McpCatalog(Map.of(), List.of()); + var executor = Executors.newVirtualThreadPerTaskExecutor()) { + var tasks = IntStream.range(0, clientCount) + .mapToObj(index -> executor.submit(() -> { + ready.countDown(); + try { + assertTrue(start.await(5, SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("interrupted", e); + } + catalog.addRemoteClient(new TestRemoteClient("remote-" + index)); + })) + .toList(); + + assertTrue(ready.await(5, SECONDS)); + start.countDown(); + for (var task : tasks) { + task.get(); + } + + assertEquals(clientCount, catalog.remoteClients().size()); + assertThrows( + UnsupportedOperationException.class, + () -> catalog.remoteClients().clear()); + } + } + + @Test + void dynamicallyAddedRemoteUsesTheNegotiatedProtocolVersion() { + var observedVersion = new AtomicReference(); + var remote = new TestRemoteClient("dynamic") { + @Override + public List listTools() { + observedVersion.set(protocolVersion()); + return List.of(); + } + }; + + try (var engine = McpEngine.builder().build()) { + var initialize = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2024_11_05.identifier()), + "capabilities", + Document.of(Map.of()), + "clientInfo", + Document.of(Map.of())))) + .build(); + assertNull(engine.execute(initialize, KnownProtocolVersion.V2024_11_05).getError()); + + engine.addRemoteClient(remote); + + assertEquals(KnownProtocolVersion.V2024_11_05, observedVersion.get()); + } + } + + @Test + void notificationRefreshDoesNotRunOnTheNotifyingThread() throws Exception { + var refreshEntered = new CountDownLatch(1); + var releaseRefresh = new CountDownLatch(1); + var calls = new AtomicInteger(); + var remote = new TestRemoteClient("notifying") { + @Override + public List listTools() { + if (calls.incrementAndGet() > 1) { + refreshEntered.countDown(); + try { + assertTrue(releaseRefresh.await(5, SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("interrupted", e); + } + } + return List.of(tool("notifying-tool")); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + try { + catalog.bindTransport(ignored -> {}, ignored -> {}); + catalog.initializeRemoteClients( + initializeRequest(), + BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25)); + + var notification = JsonRpcRequest.builder() + .jsonrpc("2.0") + .method(McpMethod.Standard.NOTIFICATIONS_TOOLS_LIST_CHANGED.wireName()) + .build(); + assertTimeoutPreemptively( + Duration.ofMillis(500), + () -> remote.sendNotification(notification)); + assertTrue(refreshEntered.await(2, SECONDS)); + } finally { + releaseRefresh.countDown(); + } + } + } + + private TestRemoteClient blockingToolClient( + String name, + String toolName, + CountDownLatch entered, + AtomicBoolean timedOut + ) { + return new TestRemoteClient(name) { + @Override + public List listTools() { + entered.countDown(); + try { + if (!entered.await(2, SECONDS)) { + timedOut.set(true); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("interrupted", e); + } + return List.of(tool(toolName)); + } + }; + } + + private ToolInfo tool(String name) { + return ToolInfo.builder() + .name(name) + .inputSchema(JsonObjectSchema.builder().build()) + .build(); + } + + private JsonRpcRequest initializeRequest() { + return JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier()), + "capabilities", + Document.of(Map.of()), + "clientInfo", + Document.of(Map.of())))) + .build(); + } + + private JsonRpcRequest request(int id, String method) { + return JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(id)) + .method(method) + .build(); + } + + private static class TestRemoteClient extends McpRemoteClient { + private final String name; + + TestRemoteClient(String name) { + this.name = name; + } + + @Override + public List listTools() { + return List.of(); + } + + @Override + public List listPrompts() { + return List.of(); + } + + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + return request.getId() == null + ? null + : JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .result(Document.of(Map.of())) + .build(); + } + + @Override + public void start() {} + + @Override + public void close() {} + + @Override + public String name() { + return name; + } + + void sendNotification(JsonRpcRequest notification) { + notify(notification); + } + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpHttpHandlerTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpHttpHandlerTest.java new file mode 100644 index 0000000000..ab3991e5b4 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpHttpHandlerTest.java @@ -0,0 +1,186 @@ +/* + * 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.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; + +class McpHttpHandlerTest { + + @Test + void loopbackHandlerAcceptsLoopbackHostAndOrigin() { + var response = handler().handle( + initializeRequest(), + Map.of( + "Host", + List.of("127.0.0.1:8080"), + "Origin", + List.of("http://localhost:3000"))); + + assertEquals(200, response.statusCode()); + assertNull(response.body().getError()); + } + + @Test + void loopbackHandlerAcceptsIpv6LoopbackHostAndOrigin() { + var response = handler().handle( + initializeRequest(), + Map.of( + "Host", + List.of("[::1]:8080"), + "Origin", + List.of("http://[::1]:3000"))); + + assertEquals(200, response.statusCode()); + assertNull(response.body().getError()); + } + + @Test + void loopbackHandlerRejectsNonLoopbackHost() { + var response = handler().handle( + initializeRequest(), + Map.of( + "Host", + List.of("evil.example.com"), + "Origin", + List.of("http://evil.example.com"))); + + assertEquals(400, response.statusCode()); + assertEquals(-32020, response.body().getError().getCode()); + } + + @Test + void loopbackHandlerRejectsNonLoopbackOrigin() { + var response = handler().handle( + initializeRequest(), + Map.of( + "Host", + List.of("localhost:8080"), + "Origin", + List.of("http://evil.example.com"))); + + assertEquals(400, response.statusCode()); + assertEquals(-32020, response.body().getError().getCode()); + } + + @Test + void protocolNegotiationIsScopedToEachHttpRequest() { + try (var engine = McpEngine.builder().build()) { + var handler = new McpHttpHandler(engine); + + var legacyInitialize = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2024_11_05.identifier())))) + .build(); + assertNull(handler.handle(legacyInitialize, Map.of()).body().getError()); + + var statelessParams = Document.of(Map.of( + "_meta", + Document.of(Map.of( + McpWireNames.PROTOCOL_VERSION, + Document.of(KnownProtocolVersion.V2026_07_28.identifier()), + McpWireNames.CLIENT_CAPABILITIES, + Document.of(Map.of()))))); + var discover = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(2)) + .method(McpMethod.Standard.SERVER_DISCOVER.wireName()) + .params(statelessParams) + .build(); + var discoverResponse = handler.handle( + discover, + Map.of( + "MCP-Protocol-Version", + List.of(KnownProtocolVersion.V2026_07_28.identifier()), + "Mcp-Method", + List.of(McpMethod.Standard.SERVER_DISCOVER.wireName()))); + assertNull(discoverResponse.body().getError()); + + var headerlessPing = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(3)) + .method(McpMethod.Standard.PING.wireName()) + .build(); + var pingResponse = handler.handle(headerlessPing, Map.of()); + + assertEquals(200, pingResponse.statusCode()); + assertNull(pingResponse.body().getError()); + } + } + + @Test + void nonStringInitializeVersionReturnsInvalidParams() { + try (var engine = McpEngine.builder().build()) { + var handler = new McpHttpHandler(engine); + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of("protocolVersion", Document.of(1)))) + .build(); + + var response = handler.handle(request, Map.of()); + + assertEquals(200, response.statusCode()); + assertEquals(-32602, response.body().getError().getCode()); + } + } + + @Test + void unsupportedLegacyProtocolHeaderReturnsBadRequest() { + try (var engine = McpEngine.builder().build()) { + var handler = new McpHttpHandler(engine); + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.PING.wireName()) + .build(); + + var response = handler.handle( + request, + Map.of("MCP-Protocol-Version", List.of("unsupported-version"))); + + assertEquals(400, response.statusCode()); + assertEquals(-32022, response.body().getError().getCode()); + } + } + + @Test + void literalBase64MarkerIsEscaped() { + var value = "=?base64?not-encoded?="; + var encoded = McpHttpBinding.encodeParameter(value); + + assertNotEquals(value, encoded); + assertEquals(value, McpHttpBinding.decodeParameter(encoded)); + } + + private McpHttpHandler handler() { + var service = McpEngine.builder().services(Map.of()).build(); + return McpHttpHandler.forLoopback(service); + } + + private JsonRpcRequest initializeRequest() { + return JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method("initialize") + .params(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier())))) + .build(); + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpSchemaFactoryTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpSchemaFactoryTest.java new file mode 100644 index 0000000000..37f67c8e0f --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpSchemaFactoryTest.java @@ -0,0 +1,147 @@ +/* + * 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 java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.schema.ApiOperation; +import software.amazon.smithy.java.core.schema.ApiService; +import software.amazon.smithy.java.core.schema.Schema; +import software.amazon.smithy.java.core.schema.SchemaIndex; +import software.amazon.smithy.java.core.schema.SerializableStruct; +import software.amazon.smithy.java.core.schema.ShapeBuilder; +import software.amazon.smithy.java.core.schema.Unit; +import software.amazon.smithy.java.core.serde.TypeRegistry; +import software.amazon.smithy.java.server.Operation; +import software.amazon.smithy.java.server.Service; +import software.amazon.smithy.model.shapes.ShapeId; + +class McpSchemaFactoryTest { + + @Test + void toolUsesRuntimeOperationNameInsteadOfSourceSchemaName() { + var service = new TestService(); + var operation = Operation.of( + "Echo", + (input, ignored) -> input, + new TestApiOperation("EchoProxy"), + service); + + var descriptor = new McpSchemaFactory(service.schemaIndex()) + .createTool("test", service, operation); + + assertEquals("Echo", descriptor.info().getName()); + assertEquals( + "This tool invokes Echo API of TestService.", + descriptor.info().getDescription()); + } + + @Test + void ordinaryToolAlsoUsesItsRuntimeOperationName() { + var service = new TestService(); + var operation = Operation.of( + "TestSimpleText", + (input, ignored) -> input, + new TestApiOperation("SimpleTextSchema"), + service); + + var descriptor = new McpSchemaFactory(service.schemaIndex()) + .createTool("test", service, operation); + + assertEquals("TestSimpleText", descriptor.info().getName()); + } + + private static final class TestService implements Service { + private static final Schema SCHEMA = + Schema.createService(ShapeId.from("example#TestService")); + + @Override + public Operation getOperation( + String operationName + ) { + return null; + } + + @Override + public List> getAllOperations() { + return List.of(); + } + + @Override + public Schema schema() { + return SCHEMA; + } + + @Override + public TypeRegistry typeRegistry() { + return TypeRegistry.empty(); + } + + @Override + public SchemaIndex schemaIndex() { + return SchemaIndex.compose(); + } + } + + private static final class TestApiOperation implements ApiOperation { + private static final ApiService SERVICE = + () -> Schema.createService(ShapeId.from("example#TestService")); + private final Schema schema; + + private TestApiOperation(String name) { + schema = Schema.createOperation(ShapeId.from("example#" + name)); + } + + @Override + public ShapeBuilder inputBuilder() { + return Unit.builder(); + } + + @Override + public ShapeBuilder outputBuilder() { + return Unit.builder(); + } + + @Override + public Schema schema() { + return schema; + } + + @Override + public Schema inputSchema() { + return Unit.SCHEMA; + } + + @Override + public Schema outputSchema() { + return Unit.SCHEMA; + } + + @Override + public TypeRegistry errorRegistry() { + return TypeRegistry.empty(); + } + + @Override + public List effectiveAuthSchemes() { + return List.of(); + } + + @Override + public List errorSchemas() { + return List.of(); + } + + @Override + public ApiService service() { + return SERVICE; + } + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/ProtocolVersionTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/ProtocolVersionTest.java index d95017e472..7f2f3e2367 100644 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/ProtocolVersionTest.java +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/ProtocolVersionTest.java @@ -6,7 +6,6 @@ package software.amazon.smithy.java.mcp.server; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; @@ -15,42 +14,43 @@ class ProtocolVersionTest { @Test void knownVersionsResolveCorrectly() { - assertInstanceOf(ProtocolVersion.v2024_11_05.class, ProtocolVersion.version("2024-11-05")); - assertInstanceOf(ProtocolVersion.v2025_03_26.class, ProtocolVersion.version("2025-03-26")); - assertInstanceOf(ProtocolVersion.v2025_06_18.class, ProtocolVersion.version("2025-06-18")); - assertInstanceOf(ProtocolVersion.v2025_11_25.class, ProtocolVersion.version("2025-11-25")); + assertEquals(KnownProtocolVersion.V2024_11_05, ProtocolVersion.parse("2024-11-05")); + assertEquals(KnownProtocolVersion.V2025_03_26, ProtocolVersion.parse("2025-03-26")); + assertEquals(KnownProtocolVersion.V2025_06_18, ProtocolVersion.parse("2025-06-18")); + assertEquals(KnownProtocolVersion.V2025_11_25, ProtocolVersion.parse("2025-11-25")); + assertEquals(KnownProtocolVersion.V2026_07_28, ProtocolVersion.parse("2026-07-28")); } @Test void unknownVersionReturnsUnknownVersion() { - var version = ProtocolVersion.version("9999-01-01"); - assertInstanceOf(ProtocolVersion.UnknownVersion.class, version); + var version = ProtocolVersion.parse("9999-01-01"); + assertTrue(version instanceof UnknownProtocolVersion); assertEquals("9999-01-01", version.identifier()); } @Test void nullVersionResolvesToDefault() { - var version = ProtocolVersion.version(null); + var version = ProtocolVersion.parse(null); assertEquals(ProtocolVersion.defaultVersion(), version); } @Test - void defaultVersionIs2025_03_26() { + void defaultVersionIsLegacyHttpCompatibilityVersion() { assertEquals("2025-03-26", ProtocolVersion.defaultVersion().identifier()); } @Test void compareToOrdersChronologically() { - assertTrue(ProtocolVersion.v2024_11_05.INSTANCE.compareTo(ProtocolVersion.v2025_03_26.INSTANCE) < 0); - assertTrue(ProtocolVersion.v2025_03_26.INSTANCE.compareTo(ProtocolVersion.v2025_06_18.INSTANCE) < 0); - assertTrue(ProtocolVersion.v2025_06_18.INSTANCE.compareTo(ProtocolVersion.v2025_11_25.INSTANCE) < 0); - assertEquals(0, ProtocolVersion.v2025_11_25.INSTANCE.compareTo(ProtocolVersion.v2025_11_25.INSTANCE)); + assertTrue(KnownProtocolVersion.V2024_11_05.compareTo(KnownProtocolVersion.V2025_03_26) < 0); + assertTrue(KnownProtocolVersion.V2025_03_26.compareTo(KnownProtocolVersion.V2025_06_18) < 0); + assertTrue(KnownProtocolVersion.V2025_06_18.compareTo(KnownProtocolVersion.V2025_11_25) < 0); + assertTrue(KnownProtocolVersion.V2025_11_25.compareTo(KnownProtocolVersion.V2026_07_28) < 0); + assertEquals(0, KnownProtocolVersion.V2026_07_28.compareTo(KnownProtocolVersion.V2026_07_28)); } @Test void knownVersionsRankAboveUnknown() { - var unknown = ProtocolVersion.version("0000-00-00"); - assertTrue(ProtocolVersion.v2024_11_05.INSTANCE.compareTo(unknown) > 0); - assertTrue(ProtocolVersion.v2025_11_25.INSTANCE.compareTo(unknown) > 0); + var unknown = ProtocolVersion.parse("0000-00-00"); + assertTrue(unknown instanceof UnknownProtocolVersion); } } diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpClientTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpClientTest.java new file mode 100644 index 0000000000..72a1383303 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpClientTest.java @@ -0,0 +1,31 @@ +/* + * 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.assertNotEquals; + +import java.math.BigInteger; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.serde.document.Document; + +class StdioMcpClientTest { + + @Test + void requestKeysSupportEveryValidNumericId() { + assertEquals("number:2147483648", StdioMcpClient.requestKey(Document.of(2_147_483_648L))); + assertEquals( + "number:9223372036854775808", + StdioMcpClient.requestKey(Document.of(new BigInteger("9223372036854775808")))); + } + + @Test + void numericAndStringIdsDoNotCollide() { + assertNotEquals( + StdioMcpClient.requestKey(Document.of(1)), + StdioMcpClient.requestKey(Document.of("1"))); + } +} 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/StdioMcpServerTest.java similarity index 83% rename from mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java rename to mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpServerTest.java index 88c7abfe6f..b97bfd24d8 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/StdioMcpServerTest.java @@ -21,7 +21,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.AfterEach; @@ -45,7 +44,7 @@ import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.model.shapes.ShapeType; -public class McpServerTest { +public class StdioMcpServerTest { private static final JsonCodec CODEC = JsonCodec.builder() .settings(JsonSettings.builder() .serializeTypeInDocuments(false) @@ -64,6 +63,22 @@ public void beforeEach() { output = new TestOutputStream(); } + @Test + void malformedJsonReturnsParseError() { + server = StdioMcpServer.builder() + .engine(McpEngine.builder().build()) + .input(input) + .output(output) + .build(); + server.start(); + + input.write("{not-json}\n"); + var response = read(); + + assertEquals(-32700, response.getError().getCode()); + assertNull(response.getId()); + } + @AfterEach public void afterEach() { if (server != null) { @@ -76,7 +91,7 @@ private void initializeWithProtocolVersion(ProtocolVersion protocolVersion) { final String expectedPv; if (protocolVersion == null) { pvDoc = Document.of(Map.of()); - expectedPv = ProtocolVersion.v2024_11_05.INSTANCE.identifier(); + expectedPv = ProtocolVersion.defaultVersion().identifier(); } else { pvDoc = Document.of(Map.of("protocolVersion", Document.of(protocolVersion.identifier()))); expectedPv = protocolVersion.identifier(); @@ -86,9 +101,26 @@ private void initializeWithProtocolVersion(ProtocolVersion protocolVersion) { assertEquals(expectedPv, pv); } + private Document modernParams(Map members) { + var params = new HashMap<>(members); + params.put("_meta", + Document.of(Map.of( + "io.modelcontextprotocol/protocolVersion", + Document.of(KnownProtocolVersion.V2026_07_28.identifier()), + "io.modelcontextprotocol/clientInfo", + Document.of(Map.of( + "name", + Document.of("test-client"), + "version", + Document.of("1.0.0"))), + "io.modelcontextprotocol/clientCapabilities", + Document.of(Map.of())))); + return Document.of(params); + } + @Test public void testPing() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -118,7 +150,7 @@ public void testPing() { @Test public void initializeWithV2025_11_25ProtocolVersion() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -132,7 +164,7 @@ public void initializeWithV2025_11_25ProtocolVersion() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_11_25.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_11_25); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -154,9 +186,130 @@ public void initializeWithV2025_11_25ProtocolVersion() { assertTrue(outputSchema.get("properties").asStringMap().containsKey("outputStr")); } + @Test + public void supportsLoggingAndCompletionUtilities() { + server = StdioMcpServer.builder() + .name("smithy-mcp-server") + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test#TestService")) + .proxyEndpoint("http://localhost") + .model(MODEL) + .build()) + .build(); + + server.start(); + + write("initialize", + Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier())))); + var capabilities = read().getResult().getMember("capabilities"); + assertNotNull(capabilities.getMember("completions")); + assertNotNull(capabilities.getMember("logging")); + + write("logging/setLevel", Document.of(Map.of("level", Document.of("info")))); + assertTrue(read().getResult().asStringMap().isEmpty()); + + write("completion/complete", + Document.of(Map.of( + "ref", + Document.of(Map.of( + "type", + Document.of("ref/prompt"), + "name", + Document.of("test_prompt"))), + "argument", + Document.of(Map.of( + "name", + Document.of("value"), + "value", + Document.of("par")))))); + var completion = read().getResult().getMember("completion"); + assertTrue(completion.getMember("values").asList().isEmpty()); + assertEquals(0, completion.getMember("total").asNumber().intValue()); + assertFalse(completion.getMember("hasMore").asBoolean()); + } + + @Test + public void supportsV2026_07_28StatelessProtocol() { + server = StdioMcpServer.builder() + .name("smithy-mcp-server") + .version("1.2.3") + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test#TestService")) + .proxyEndpoint("http://localhost") + .model(MODEL) + .build()) + .build(); + + server.start(); + + write("server/discover", modernParams(Map.of())); + var discoverResponse = read(); + assertNull( + discoverResponse.getError(), + () -> discoverResponse.getError() == null ? null : discoverResponse.getError().getMessage()); + var discover = discoverResponse.getResult(); + assertEquals("complete", discover.getMember("resultType").asString()); + assertEquals(0, discover.getMember("ttlMs").asNumber().intValue()); + assertEquals("private", discover.getMember("cacheScope").asString()); + assertEquals( + "2026-07-28", + discover.getMember("supportedVersions").asList().getFirst().asString()); + assertNotNull(discover.getMember("capabilities").getMember("tools")); + assertNotNull(discover.getMember("capabilities").getMember("prompts")); + assertEquals( + "smithy-mcp-server", + discover.getMember("_meta") + .getMember("io.modelcontextprotocol/serverInfo") + .getMember("name") + .asString()); + + write("tools/list", modernParams(Map.of())); + var tools = read().getResult(); + assertEquals("complete", tools.getMember("resultType").asString()); + assertEquals(0, tools.getMember("ttlMs").asNumber().intValue()); + assertEquals("private", tools.getMember("cacheScope").asString()); + assertEquals(6, tools.getMember("tools").asList().size()); + + write("ping", modernParams(Map.of())); + assertEquals(-32601, read().getError().getCode()); + } + + @Test + public void modernProtocolRejectsMissingMetadata() { + server = StdioMcpServer.builder() + .name("smithy-mcp-server") + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test#TestService")) + .proxyEndpoint("http://localhost") + .model(MODEL) + .build()) + .build(); + + server.start(); + + var incompleteMeta = Document.of(Map.of( + "_meta", + Document.of(Map.of( + "io.modelcontextprotocol/protocolVersion", + Document.of("2026-07-28"))))); + write("server/discover", incompleteMeta); + assertEquals(-32602, read().getError().getCode()); + } + @Test public void noOutputSchemaWithUnsupportedProtocolVersion() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -170,7 +323,7 @@ public void noOutputSchemaWithUnsupportedProtocolVersion() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_03_26.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_03_26); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -188,7 +341,7 @@ public void noOutputSchemaWithUnsupportedProtocolVersion() { @Test public void validateToolsList() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -224,7 +377,7 @@ public void validateToolsList() { @Test public void validateNoIOOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -238,7 +391,7 @@ public void validateNoIOOperationTool() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -258,7 +411,7 @@ public void validateNoIOOperationTool() { @Test public void validateNoOutputOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -272,7 +425,7 @@ public void validateNoOutputOperationTool() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -301,7 +454,7 @@ public void validateNoOutputOperationTool() { @Test public void validateNoInputOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -315,7 +468,7 @@ public void validateNoInputOperationTool() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -344,7 +497,7 @@ public void validateNoInputOperationTool() { @Test public void validateTestOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -358,7 +511,7 @@ public void validateTestOperationTool() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -378,7 +531,7 @@ public void validateTestOperationTool() { @Test void readOnlyOperationHasReadOnlyHintAnnotation() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -410,7 +563,7 @@ void readOnlyOperationHasReadOnlyHintAnnotation() { @Test void idempotentOperationHasIdempotentHintAnnotation() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -442,7 +595,7 @@ void idempotentOperationHasIdempotentHintAnnotation() { @Test void plainOperationHasNoAnnotations() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -472,7 +625,7 @@ void plainOperationHasNoAnnotations() { @Test void annotationsStrippedForOldProtocolVersion() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -486,7 +639,7 @@ void annotationsStrippedForOldProtocolVersion() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2024_11_05.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2024_11_05); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -502,7 +655,7 @@ void annotationsStrippedForOldProtocolVersion() { @Test void testNumberAndStringIds() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -522,6 +675,20 @@ void testNumberAndStringIds() { assertEquals(42, response.getId().asNumber().intValue()); assertNotNull(response.getResult()); + // Test with long ID + var longId = (long) Integer.MAX_VALUE + 1; + write("tools/list", Document.of(Map.of()), Document.of(longId)); + response = read(); + assertEquals(longId, response.getId().asLong()); + assertNotNull(response.getResult()); + + // Test with arbitrary precision integer ID + var bigIntegerId = BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.ONE); + write("tools/list", Document.of(Map.of()), Document.of(bigIntegerId)); + response = read(); + assertEquals(bigIntegerId, response.getId().asBigInteger()); + assertNotNull(response.getResult()); + // Test with string ID write("tools/list", Document.of(Map.of()), Document.of("test-id-1")); response = read(); @@ -548,7 +715,7 @@ void testNumberAndStringIds() { @Test void testInvalidIds() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -591,7 +758,7 @@ void testInvalidIds() { @Test void testRequestsRequireIds() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -605,17 +772,16 @@ void testRequestsRequireIds() { server.start(); - // Test regular request without ID (should fail with specific message) + // JSON-RPC treats any request without an ID as a notification, even when the + // method name is not in the notifications namespace. write("tools/list", Document.of(Map.of()), null); - var response = read(); - assertNotNull(response.getError()); - assertTrue(response.getError().getMessage().contains("Requests are expected to have ids")); + output.assertNoOutput(); } @Test void testInputAdaptation() { AtomicReference capturedInput = new AtomicReference<>(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -726,7 +892,7 @@ public void readBeforeSerialization(InputHook hook) { @Test void testNotificationsDoNotRequireRequestId() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -756,7 +922,7 @@ void testNotificationsDoNotRequireRequestId() { @Test void testUnknownMethodReturnsMethodNotFound() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -796,7 +962,7 @@ void testUnknownMethodReturnsMethodNotFound() { @Test void testUnknownNotificationIsSilentlyDropped() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -827,7 +993,7 @@ void testUnknownNotificationIsSilentlyDropped() { @Test void testPromptsList() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -871,7 +1037,7 @@ void testPromptsList() { @Test void testPromptsGetWithValidPrompt() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -905,7 +1071,7 @@ void testPromptsGetWithValidPrompt() { @Test void testPromptsGetWithDifferentCasing() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -976,7 +1142,7 @@ void testPromptsGetWithDifferentCasing() { @Test void testPromptsGetWithInvalidPrompt() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1016,7 +1182,7 @@ void testPromptsGetWithTemplateArguments() { .assemble() .unwrap(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1058,7 +1224,7 @@ void testPromptsGetWithMissingRequiredArguments() { .assemble() .unwrap(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1096,7 +1262,7 @@ void testApplyTemplateArgumentsEdgeCases() { .assemble() .unwrap(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1581,7 +1747,7 @@ private void writeNotification(String method, Document params) { @Test void testUnionSchemaGeneratesOneOfWithWrappedMembers() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1595,7 +1761,7 @@ void testUnionSchemaGeneratesOneOfWithWrappedMembers() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -1636,7 +1802,7 @@ void testUnionSchemaGeneratesOneOfWithWrappedMembers() { @Test void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1650,7 +1816,7 @@ void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { server.start(); - initializeWithProtocolVersion(ProtocolVersion.v2025_06_18.INSTANCE); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_06_18); write("tools/list", Document.of(Map.of())); var response = read(); var tools = response.getResult().asStringMap().get("tools").asList(); @@ -1676,13 +1842,13 @@ void testToolsListChangedNotificationInvalidatesCache() { var callCounter = new AtomicInteger(0); var mockProxy = new CacheTestProxy(callCounter); - var service = McpService.builder() + var service = McpEngine.builder() .name("test") - .proxyList(List.of(mockProxy)) + .remoteClients(List.of(mockProxy)) .build(); var notifications = new ArrayList(); - service.setNotificationWriter(notifications::add); + service.bindTransport(notifications::add, ignored -> {}); // Initialize to set up proxies var initRequest = JsonRpcRequest.builder() @@ -1691,7 +1857,7 @@ void testToolsListChangedNotificationInvalidatesCache() { .params(Document.of(Map.of("protocolVersion", Document.of("2024-11-05")))) .jsonrpc("2.0") .build(); - service.handleRequest(initRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(initRequest, ProtocolVersion.defaultVersion()); // Verify notifications/initialized was sent during initialization assertTrue(mockProxy.getSentNotifications().contains("notifications/initialized"), @@ -1704,11 +1870,11 @@ void testToolsListChangedNotificationInvalidatesCache() { .params(Document.of(Map.of())) .jsonrpc("2.0") .build(); - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); assertEquals(1, callCounter.get(), "First call should fetch from proxy"); // Second tools/list - uses cache - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); assertEquals(1, callCounter.get(), "Second call should use cache"); // Send tools/list_changed notification @@ -1723,12 +1889,18 @@ 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"); + assertTimeoutPreemptively(Duration.ofSeconds(2), () -> { + while (callCounter.get() < 2) { + Thread.sleep(10); + } + }); + + // Third tools/list - should use the asynchronously refreshed cache + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); + assertEquals(2, callCounter.get(), "Notification should refresh before the third call"); // Fourth tools/list - uses cache again (counter should NOT increment) - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); assertEquals(2, callCounter.get(), "Fourth call should use cache (not increment to 3)"); } @@ -1737,13 +1909,13 @@ void testOtherNotificationsDoNotInvalidateCache() { var callCounter = new AtomicInteger(0); var mockProxy = new CacheTestProxy(callCounter); - var service = McpService.builder() + var service = McpEngine.builder() .name("test") - .proxyList(List.of(mockProxy)) + .remoteClients(List.of(mockProxy)) .build(); var notifications = new ArrayList(); - service.setNotificationWriter(notifications::add); + service.bindTransport(notifications::add, ignored -> {}); // Initialize var initRequest = JsonRpcRequest.builder() @@ -1752,7 +1924,7 @@ void testOtherNotificationsDoNotInvalidateCache() { .params(Document.of(Map.of("protocolVersion", Document.of("2024-11-05")))) .jsonrpc("2.0") .build(); - service.handleRequest(initRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(initRequest, ProtocolVersion.defaultVersion()); // Verify notifications/initialized was sent during initialization assertTrue(mockProxy.getSentNotifications().contains("notifications/initialized"), @@ -1765,7 +1937,7 @@ void testOtherNotificationsDoNotInvalidateCache() { .params(Document.of(Map.of())) .jsonrpc("2.0") .build(); - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); assertEquals(1, callCounter.get()); // Send different notification @@ -1780,11 +1952,11 @@ void testOtherNotificationsDoNotInvalidateCache() { assertEquals("notifications/prompts/list_changed", notifications.get(0).getMethod()); // Second tools/list - should still use cache - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); assertEquals(1, callCounter.get(), "Cache should not be invalidated by other notifications"); } - private static class CacheTestProxy extends McpServerProxy { + private static class CacheTestProxy extends McpRemoteClient { private final AtomicInteger callCounter; private final List sentNotifications = new ArrayList<>(); @@ -1809,18 +1981,17 @@ public List listPrompts() { } @Override - protected CompletableFuture rpc(JsonRpcRequest request) { + protected JsonRpcResponse exchange(JsonRpcRequest request) { // Notifications have no ID if (request.getId() == null) { sentNotifications.add(request.getMethod()); - return CompletableFuture.completedFuture(null); + return null; } - return CompletableFuture.completedFuture( - JsonRpcResponse.builder() - .id(request.getId()) - .result(Document.of(Map.of())) - .jsonrpc("2.0") - .build()); + return JsonRpcResponse.builder() + .id(request.getId()) + .result(Document.of(Map.of())) + .jsonrpc("2.0") + .build(); } List getSentNotifications() { @@ -1828,12 +1999,10 @@ List getSentNotifications() { } @Override - protected void start() {} + public void start() {} @Override - protected CompletableFuture shutdown() { - return CompletableFuture.completedFuture(null); - } + public void close() {} @Override public String name() { @@ -1850,9 +2019,9 @@ void sendNotification(JsonRpcRequest notification) { @Test void testReadBeforeAndAfterExecution() { var capturedMethod = new AtomicReference(); - var capturedResponse = new AtomicReference(); + var capturedResponse = new AtomicReference(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1862,16 +2031,16 @@ void testReadBeforeAndAfterExecution() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { - capturedMethod.set(hook.request().getMethod()); + public void readBeforeExecution(McpExecutionContext hook) { + capturedMethod.set(hook.call().method().wireName()); } @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { capturedResponse.set(response); @@ -1894,7 +2063,7 @@ void testReadBeforeAndAfterToolCallLocal() { var capturedIsProxy = new AtomicReference(); var afterToolCallFired = new AtomicReference<>(false); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1904,18 +2073,18 @@ void testReadBeforeAndAfterToolCallLocal() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeToolCall(McpToolCallHook hook) { - capturedToolName.set(hook.toolName()); + public void readBeforeToolCall(McpToolExecutionContext hook) { + capturedToolName.set(hook.call().name()); capturedServerId.set(hook.serverId()); - capturedIsProxy.set(hook.isProxy()); + capturedIsProxy.set(hook.remote()); } @Override public void readAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, + McpToolExecutionContext hook, + McpOutcome response, RuntimeException error ) { afterToolCallFired.set(true); @@ -1947,7 +2116,7 @@ void testReadBeforeAndAfterToolCallProxy() { var afterToolCallFired = new AtomicReference<>(false); var mockProxy = new CacheTestProxy(new AtomicInteger(0)); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1957,18 +2126,18 @@ void testReadBeforeAndAfterToolCallProxy() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .addService(mockProxy) - .interceptor(new McpServerInterceptor() { + .addRemoteClient(mockProxy) + .interceptor(new McpInterceptor() { @Override - public void readBeforeToolCall(McpToolCallHook hook) { - capturedToolName.set(hook.toolName()); - capturedIsProxy.set(hook.isProxy()); + public void readBeforeToolCall(McpToolExecutionContext hook) { + capturedToolName.set(hook.call().name()); + capturedIsProxy.set(hook.remote()); } @Override public void readAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, + McpToolExecutionContext hook, + McpOutcome response, RuntimeException error ) { afterToolCallFired.set(true); @@ -1996,7 +2165,7 @@ public void readAfterToolCall( void testReadAfterExecutionAlwaysFires() { var afterCount = new AtomicInteger(0); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2006,11 +2175,11 @@ void testReadAfterExecutionAlwaysFires() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { afterCount.incrementAndGet(); @@ -2038,7 +2207,7 @@ void testReadAfterToolCallFiresWhenBeforeToolCallThrows() { var afterToolCallFired = new AtomicReference<>(false); var capturedErrorMessage = new AtomicReference(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2048,16 +2217,16 @@ void testReadAfterToolCallFiresWhenBeforeToolCallThrows() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeToolCall(McpToolCallHook hook) { + public void readBeforeToolCall(McpToolExecutionContext hook) { throw new RuntimeException("blocked"); } @Override public void readAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, + McpToolExecutionContext hook, + McpOutcome response, RuntimeException error ) { afterToolCallFired.set(true); @@ -2077,10 +2246,12 @@ public void readAfterToolCall( Document.of("NoIOOperation"), "arguments", Document.of(Map.of())))); - read(); + var response = read(); assertTrue(afterToolCallFired.get()); assertEquals("blocked", capturedErrorMessage.get()); + assertEquals(-32603, response.getError().getCode()); + assertEquals("Internal error", response.getError().getMessage()); } @Test @@ -2088,7 +2259,7 @@ void testReadAfterExecutionFiresForProxyToolCall() { var afterExecutionFired = new AtomicReference<>(false); var mockProxy = new CacheTestProxy(new AtomicInteger(0)); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2098,12 +2269,12 @@ void testReadAfterExecutionFiresForProxyToolCall() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .addService(mockProxy) - .interceptor(new McpServerInterceptor() { + .addRemoteClient(mockProxy) + .interceptor(new McpInterceptor() { @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { afterExecutionFired.set(true); @@ -2130,7 +2301,7 @@ void testReadBeforeExecutionThrowSkipsToolHooks() { var beforeToolCallFired = new AtomicReference<>(false); var afterExecutionError = new AtomicReference(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2140,26 +2311,26 @@ void testReadBeforeExecutionThrowSkipsToolHooks() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { - if ("tools/call".equals(hook.request().getMethod())) { + public void readBeforeExecution(McpExecutionContext hook) { + if ("tools/call".equals(hook.call().method().wireName())) { throw new RuntimeException("execution-blocked"); } } @Override - public void readBeforeToolCall(McpToolCallHook hook) { + public void readBeforeToolCall(McpToolExecutionContext hook) { beforeToolCallFired.set(true); } @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { - if ("tools/call".equals(hook.request().getMethod())) { + if ("tools/call".equals(hook.call().method().wireName())) { afterExecutionError.set(error); } } @@ -2188,7 +2359,7 @@ void testContextPassesBetweenReadHooks() { Context.Key START_KEY = Context.key("start"); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2198,19 +2369,19 @@ void testContextPassesBetweenReadHooks() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { - hook.context().put(START_KEY, System.nanoTime()); + public void readBeforeExecution(McpExecutionContext hook) { + hook.requestContext().attributes().put(START_KEY, System.nanoTime()); } @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { - long start = hook.context().get(START_KEY); + long start = hook.requestContext().attributes().get(START_KEY); duration.set(System.nanoTime() - start); } }) @@ -2228,7 +2399,7 @@ public void readAfterExecution( @Test void testModifyBeforeExecutionRewritesRequest() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2238,15 +2409,10 @@ void testModifyBeforeExecutionRewritesRequest() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public JsonRpcRequest modifyBeforeExecution(McpExecutionHook hook) { - return JsonRpcRequest.builder() - .id(hook.request().getId()) - .method("ping") - .params(Document.of(Map.of())) - .jsonrpc("2.0") - .build(); + public McpCall modifyBeforeExecution(McpExecutionContext hook) { + return new McpCall.Ping(hook.call().id(), hook.call().metadata()); } }) .build(); @@ -2261,7 +2427,7 @@ public JsonRpcRequest modifyBeforeExecution(McpExecutionHook hook) { void testModifyBeforeToolCallModifiesRequest() { var modifyHookCalled = new AtomicReference<>(false); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2271,11 +2437,11 @@ void testModifyBeforeToolCallModifiesRequest() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { + public McpCall.CallTool modifyBeforeToolCall(McpToolExecutionContext hook) { modifyHookCalled.set(true); - return hook.request(); + return hook.call(); } }) .build(); @@ -2297,7 +2463,7 @@ public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { @Test void testModifyAfterExecutionTransformsResponse() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2307,21 +2473,19 @@ void testModifyAfterExecutionTransformsResponse() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public JsonRpcResponse modifyAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterExecution( + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { if (error != null) { throw error; } - return JsonRpcResponse.builder() - .id(hook.request().getId()) - .result(Document.of(Map.of("modified", Document.of("true")))) - .jsonrpc("2.0") - .build(); + return new McpOutcome.Success( + hook.call().id(), + Document.of(Map.of("modified", Document.of("true")))); } }) .build(); @@ -2335,7 +2499,7 @@ public JsonRpcResponse modifyAfterExecution( @Test void testModifyAfterToolCallTransformsResponse() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2345,19 +2509,17 @@ void testModifyAfterToolCallTransformsResponse() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public JsonRpcResponse modifyAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterToolCall( + McpToolExecutionContext hook, + McpOutcome response, RuntimeException error ) { // Always return custom response, ignoring any tool error - return JsonRpcResponse.builder() - .id(hook.request().getId()) - .result(Document.of(Map.of("tool-modified", Document.of("true")))) - .jsonrpc("2.0") - .build(); + return new McpOutcome.Success( + hook.call().id(), + Document.of(Map.of("tool-modified", Document.of("true")))); } }) .build(); @@ -2382,7 +2544,7 @@ public JsonRpcResponse modifyAfterToolCall( void testReadBeforeExecutionThrowShortCircuits() { var afterExecutionError = new AtomicReference(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2392,16 +2554,16 @@ void testReadBeforeExecutionThrowShortCircuits() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { + public void readBeforeExecution(McpExecutionContext hook) { throw new RuntimeException("blocked"); } @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { afterExecutionError.set(error); @@ -2414,14 +2576,15 @@ public void readAfterExecution( var response = read(); assertNotNull(response.getError()); - assertTrue(response.getError().getMessage().contains("blocked")); + assertEquals(-32603, response.getError().getCode()); + assertEquals("Internal error", response.getError().getMessage()); assertNotNull(afterExecutionError.get()); assertEquals("blocked", afterExecutionError.get().getMessage()); } @Test void testModifyAfterExecutionCanRecoverFromError() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2431,24 +2594,20 @@ void testModifyAfterExecutionCanRecoverFromError() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(new McpServerInterceptor() { + .interceptor(new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { + public void readBeforeExecution(McpExecutionContext hook) { throw new RuntimeException("original-error"); } @Override - public JsonRpcResponse modifyAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterExecution( + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { // Recover from the error by returning a success response - return JsonRpcResponse.builder() - .id(hook.request().getId()) - .result(Document.of(Map.of())) - .jsonrpc("2.0") - .build(); + return new McpOutcome.Success(hook.call().id(), Document.of(Map.of())); } }) .build(); @@ -2467,7 +2626,7 @@ public JsonRpcResponse modifyAfterExecution( void testChainReadHooksInvokedInOrder() { var order = new ArrayList(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2477,32 +2636,32 @@ void testChainReadHooksInvokedInOrder() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(McpServerInterceptor.chain(List.of( - new McpServerInterceptor() { + .interceptor(McpInterceptor.chain(List.of( + new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { + public void readBeforeExecution(McpExecutionContext hook) { order.add("A-before"); } @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { order.add("A-after"); } }, - new McpServerInterceptor() { + new McpInterceptor() { @Override - public void readBeforeExecution(McpExecutionHook hook) { + public void readBeforeExecution(McpExecutionContext hook) { order.add("B-before"); } @Override public void readAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { order.add("B-after"); @@ -2519,9 +2678,10 @@ public void readAfterExecution( @Test void testChainModifyBeforeToolCallPropagatesRequest() { - var capturedRequest = new AtomicReference(); + var replacement = new AtomicReference(); + var capturedRequest = new AtomicReference(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2531,26 +2691,25 @@ void testChainModifyBeforeToolCallPropagatesRequest() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(McpServerInterceptor.chain(List.of( - new McpServerInterceptor() { + .interceptor(McpInterceptor.chain(List.of( + new McpInterceptor() { @Override - public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { - var params = hook.request().getParams().asStringMap(); - var newParams = new HashMap<>(params); - newParams.put("injected", Document.of("from-first")); - return JsonRpcRequest.builder() - .id(hook.request().getId()) - .method(hook.request().getMethod()) - .params(Document.of(newParams)) - .jsonrpc(hook.request().getJsonrpc()) - .build(); + public McpCall.CallTool modifyBeforeToolCall(McpToolExecutionContext hook) { + var call = hook.call(); + var modified = new McpCall.CallTool( + call.id(), + call.name(), + call.arguments(), + call.metadata()); + replacement.set(modified); + return modified; } }, - new McpServerInterceptor() { + new McpInterceptor() { @Override - public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { - capturedRequest.set(hook.request()); - return hook.request(); + public McpCall.CallTool modifyBeforeToolCall(McpToolExecutionContext hook) { + capturedRequest.set(hook.call()); + return hook.call(); } }))) .build(); @@ -2567,9 +2726,7 @@ public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { read(); assertNotNull(capturedRequest.get()); - var injected = capturedRequest.get().getParams().getMember("injected"); - assertNotNull(injected); - assertEquals("from-first", injected.asString()); + assertTrue(replacement.get() == capturedRequest.get()); } @Test @@ -2579,7 +2736,7 @@ void testChainModifyAfterExecutionErrorPropagates() { // an error response. var secondInterceptorCalled = new AtomicReference<>(false); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2589,22 +2746,22 @@ void testChainModifyAfterExecutionErrorPropagates() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(McpServerInterceptor.chain(List.of( - new McpServerInterceptor() { + .interceptor(McpInterceptor.chain(List.of( + new McpInterceptor() { @Override - public JsonRpcResponse modifyAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterExecution( + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { throw new RuntimeException("first-error"); } }, - new McpServerInterceptor() { + new McpInterceptor() { @Override - public JsonRpcResponse modifyAfterExecution( - McpExecutionHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterExecution( + McpExecutionContext hook, + McpOutcome response, RuntimeException error ) { secondInterceptorCalled.set(true); @@ -2620,7 +2777,8 @@ public JsonRpcResponse modifyAfterExecution( // Second interceptor never runs — exception propagates immediately assertFalse(secondInterceptorCalled.get()); assertNotNull(response.getError()); - assertTrue(response.getError().getMessage().contains("first-error")); + assertEquals(-32603, response.getError().getCode()); + assertEquals("Internal error", response.getError().getMessage()); } @Test @@ -2630,7 +2788,7 @@ void testChainModifyAfterToolCallErrorPropagates() { // an error response. var secondInterceptorCalled = new AtomicReference<>(false); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2640,22 +2798,22 @@ void testChainModifyAfterToolCallErrorPropagates() { .proxyEndpoint("http://localhost") .model(MODEL) .build()) - .interceptor(McpServerInterceptor.chain(List.of( - new McpServerInterceptor() { + .interceptor(McpInterceptor.chain(List.of( + new McpInterceptor() { @Override - public JsonRpcResponse modifyAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterToolCall( + McpToolExecutionContext hook, + McpOutcome response, RuntimeException error ) { throw new RuntimeException("tool-error"); } }, - new McpServerInterceptor() { + new McpInterceptor() { @Override - public JsonRpcResponse modifyAfterToolCall( - McpToolCallHook hook, - JsonRpcResponse response, + public McpOutcome modifyAfterToolCall( + McpToolExecutionContext hook, + McpOutcome response, RuntimeException error ) { secondInterceptorCalled.set(true); @@ -2678,6 +2836,7 @@ public JsonRpcResponse modifyAfterToolCall( // Second interceptor never runs — exception propagates immediately assertFalse(secondInterceptorCalled.get()); assertNotNull(response.getError()); - assertTrue(response.getError().getMessage().contains("tool-error")); + assertEquals(-32603, response.getError().getCode()); + assertEquals("Internal error", response.getError().getMessage()); } } diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/TestInputStream.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/TestInputStream.java deleted file mode 100644 index b0802acd4d..0000000000 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/TestInputStream.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package software.amazon.smithy.java.mcp.server; - -import java.io.IOException; -import java.io.InputStream; -import java.nio.charset.StandardCharsets; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingQueue; - -final class TestInputStream extends InputStream { - private byte[] onDeck; - private int pos; - private final BlockingQueue bytes = new LinkedBlockingQueue<>(); - - void write(String s) { - bytes.add(s.getBytes(StandardCharsets.UTF_8)); - } - - void write(byte[] bytes) { - this.bytes.add(bytes); - } - - @Override - public int read() { - load(true); - return onDeck[pos++] & 0xFF; - } - - @Override - public int read(byte[] b, int off, int len) { - int rem = len; - int read = 0; - boolean first = true; - while (rem > 0) { - if (load(first) || onDeck == null) { - break; - } - first = false; - int toRead = Math.min(onDeck.length - pos, rem); - System.arraycopy(onDeck, pos, b, off, toRead); - pos += toRead; - off += toRead; - rem -= toRead; - read += toRead; - } - return read; - } - - private boolean load(boolean first) { - try { - if (onDeck == null || pos == onDeck.length) { - onDeck = first ? bytes.take() : bytes.poll(); - pos = 0; - } - return false; - } catch (Exception e) { - throw new RuntimeException(e); - } - } - - @Override - public void close() throws IOException { - super.close(); - } -} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/spi/ExternalProtocolApiTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/spi/ExternalProtocolApiTest.java new file mode 100644 index 0000000000..c169fd5f71 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/spi/ExternalProtocolApiTest.java @@ -0,0 +1,52 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.mcp.server.spi; + +import static org.junit.jupiter.api.Assertions.assertNull; + +import java.util.Set; +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.server.ExtensionMcpProtocol; +import software.amazon.smithy.java.mcp.server.McpEngine; +import software.amazon.smithy.java.mcp.server.McpMethod; +import software.amazon.smithy.java.mcp.server.McpProtocolId; +import software.amazon.smithy.java.mcp.server.UnknownProtocolVersion; + +class ExternalProtocolApiTest { + + @Test + void externalPackageCanImplementAndRegisterAProtocol() { + var protocol = new ExternalProtocol(); + try (var engine = McpEngine.builder() + .discoverProtocols(false) + .addProtocol(protocol) + .build()) { + var response = engine.execute( + JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.PING.wireName()) + .build(), + new UnknownProtocolVersion(protocol.id().identifier())); + + assertNull(response.getError()); + } + } + + private record ExternalProtocol() implements ExtensionMcpProtocol { + @Override + public McpProtocolId id() { + return McpProtocolId.of("2099-external-api"); + } + + @Override + public Set supportedMethods() { + return Set.of(McpMethod.Standard.PING); + } + } +} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/utils/TestJavaCodegenRunner.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/utils/TestJavaCodegenRunner.java index 429a929a8e..6116dfc06d 100644 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/utils/TestJavaCodegenRunner.java +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/utils/TestJavaCodegenRunner.java @@ -34,12 +34,33 @@ public static void main(String[] args) { .discoverModels(TestJavaCodegenRunner.class.getClassLoader()) .assemble() .unwrap(); + var fileManifest = FileManifest.create(Paths.get(System.getenv("output"))); + execute(plugin, + model, + fileManifest, + "smithy.java.mcp.test#TestService", + "software.amazon.smithy.java.mcp.test"); + execute( + plugin, + model, + fileManifest, + "software.amazon.smithy.java.mcp.conformance#ConformanceService", + "software.amazon.smithy.java.mcp.conformance"); + } + + private static void execute( + SmithyBuildPlugin plugin, + Model model, + FileManifest fileManifest, + String service, + String namespace + ) { PluginContext context = PluginContext.builder() - .fileManifest(FileManifest.create(Paths.get(System.getenv("output")))) + .fileManifest(fileManifest) .settings( ObjectNode.builder() - .withMember("service", "smithy.java.mcp.test#TestService") - .withMember("namespace", "software.amazon.smithy.java.mcp.test") + .withMember("service", service) + .withMember("namespace", namespace) .withMember("modes", ArrayNode.fromStrings("server")) .withMember("runtimeTraits", fromStrings("smithy.api#documentation", diff --git a/mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/TestInputStream.java b/mcp/mcp-server/src/testFixtures/java/software/amazon/smithy/java/mcp/server/TestInputStream.java similarity index 100% rename from mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/TestInputStream.java rename to mcp/mcp-server/src/testFixtures/java/software/amazon/smithy/java/mcp/server/TestInputStream.java diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java b/mcp/mcp-server/src/testFixtures/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java similarity index 100% rename from mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java rename to mcp/mcp-server/src/testFixtures/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java diff --git a/smithy-ai-traits/model/mcp.smithy b/smithy-ai-traits/model/mcp.smithy new file mode 100644 index 0000000000..eecf4846c8 --- /dev/null +++ b/smithy-ai-traits/model/mcp.smithy @@ -0,0 +1,12 @@ +$version: "2" + +namespace smithy.ai + +/// Mirrors a string member into an MCP HTTP `Mcp-Param-*` header. +/// +/// The trait value is the suffix appended to `Mcp-Param-`. Servers validate +/// that the decoded header value matches the corresponding JSON body member. +@unstable +@trait(selector: ":is(member)") +@pattern("^[A-Za-z0-9][A-Za-z0-9_-]*$") +string mcpHeader diff --git a/smithy-ai-traits/src/main/java/software/amazon/smithy/ai/McpHeaderValidator.java b/smithy-ai-traits/src/main/java/software/amazon/smithy/ai/McpHeaderValidator.java new file mode 100644 index 0000000000..1d2c203a53 --- /dev/null +++ b/smithy-ai-traits/src/main/java/software/amazon/smithy/ai/McpHeaderValidator.java @@ -0,0 +1,42 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.ai; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import software.amazon.smithy.model.Model; +import software.amazon.smithy.model.shapes.MemberShape; +import software.amazon.smithy.model.shapes.Shape; +import software.amazon.smithy.model.shapes.ShapeType; +import software.amazon.smithy.model.validation.AbstractValidator; +import software.amazon.smithy.model.validation.ValidationEvent; + +/** + * Validates that {@code mcpHeader} is only applied to members targeting strings. + */ +public final class McpHeaderValidator extends AbstractValidator { + + @Override + public List validate(Model model) { + List events = new ArrayList<>(); + for (Shape shape : model.toSet()) { + Optional trait = shape.getTrait(McpHeaderTrait.class); + if (!trait.isPresent()) { + continue; + } + + MemberShape member = shape.asMemberShape().orElseThrow(IllegalStateException::new); + Shape target = model.expectShape(member.getTarget()); + if (target.getType() != ShapeType.STRING) { + events.add(error( + member, + "The smithy.ai#mcpHeader trait can only be applied to members targeting strings.")); + } + } + return events; + } +} diff --git a/smithy-ai-traits/src/main/resources/META-INF/services/software.amazon.smithy.model.validation.Validator b/smithy-ai-traits/src/main/resources/META-INF/services/software.amazon.smithy.model.validation.Validator index 3879bd4393..a8a8acc5e2 100644 --- a/smithy-ai-traits/src/main/resources/META-INF/services/software.amazon.smithy.model.validation.Validator +++ b/smithy-ai-traits/src/main/resources/META-INF/services/software.amazon.smithy.model.validation.Validator @@ -1 +1,2 @@ software.amazon.smithy.ai.PromptUniquenessValidator +software.amazon.smithy.ai.McpHeaderValidator diff --git a/smithy-ai-traits/src/test/java/software/amazon/smithy/ai/McpHeaderValidatorTest.java b/smithy-ai-traits/src/test/java/software/amazon/smithy/ai/McpHeaderValidatorTest.java new file mode 100644 index 0000000000..87a4b75ad5 --- /dev/null +++ b/smithy-ai-traits/src/test/java/software/amazon/smithy/ai/McpHeaderValidatorTest.java @@ -0,0 +1,45 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.ai; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.Objects; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.model.Model; +import software.amazon.smithy.model.validation.ValidatedResult; + +class McpHeaderValidatorTest { + + @Test + void acceptsStringMembers() { + var result = assemble("/mcp-header-valid.smithy"); + + assertEquals(0, + result.getValidationEvents() + .stream() + .filter(event -> event.getMessage().contains("mcpHeader trait can only")) + .count()); + } + + @Test + void rejectsNonStringMembers() { + var result = assemble("/mcp-header-invalid.smithy"); + + assertEquals(1, + result.getValidationEvents() + .stream() + .filter(event -> event.getMessage().contains("mcpHeader trait can only")) + .count()); + } + + private ValidatedResult assemble(String resource) { + return Model.assembler() + .addImport(Objects.requireNonNull(getClass().getResource(resource))) + .discoverModels(getClass().getClassLoader()) + .assemble(); + } +} diff --git a/smithy-ai-traits/src/test/resources/mcp-header-invalid.smithy b/smithy-ai-traits/src/test/resources/mcp-header-invalid.smithy new file mode 100644 index 0000000000..5d1b382bbf --- /dev/null +++ b/smithy-ai-traits/src/test/resources/mcp-header-invalid.smithy @@ -0,0 +1,10 @@ +$version: "2" + +namespace smithy.ai.test + +use smithy.ai#mcpHeader + +structure InvalidMcpHeaderInput { + @mcpHeader("tenant") + value: Integer +} diff --git a/smithy-ai-traits/src/test/resources/mcp-header-valid.smithy b/smithy-ai-traits/src/test/resources/mcp-header-valid.smithy new file mode 100644 index 0000000000..fff6c6767b --- /dev/null +++ b/smithy-ai-traits/src/test/resources/mcp-header-valid.smithy @@ -0,0 +1,10 @@ +$version: "2" + +namespace smithy.ai.test + +use smithy.ai#mcpHeader + +structure ValidMcpHeaderInput { + @mcpHeader("tenant") + value: String +}