From 407a09df35c329f3fda4f92975ef4210e2d40a35 Mon Sep 17 00:00:00 2001 From: Adwait Kumar Singh Date: Sat, 29 Aug 2026 06:43:58 +0530 Subject: [PATCH] Implement modern MCP protocol conformance --- examples/mcp-server/README.md | 57 +- examples/mcp-server/build.gradle.kts | 6 +- .../example/server/mcp/MCPServerExample.java | 8 +- .../example/server/mcp/ProxyMCPExample.java | 8 +- mcp/mcp-schemas/model/main.smithy | 5 + mcp/mcp-server/README.md | 221 +++ mcp/mcp-server/build.gradle.kts | 18 + .../java/mcp/server/McpConformanceTest.java | 302 ++++ ...ava => StdioMcpServerIntegrationTest.java} | 137 +- .../java/mcp/server/TestOutputStream.java | 88 - .../META-INF/smithy/conformance.smithy | 82 + .../it/resources/META-INF/smithy/main.smithy | 2 +- .../src/it/resources/META-INF/smithy/manifest | 3 +- .../conformance-baseline-2025-11-25.yaml | 22 + .../conformance-baseline-2026-07-28.yaml | 37 + .../java/mcp/server/BuiltInProtocol.java | 78 + .../java/mcp/server/BuiltInProtocols.java | 83 + .../java/mcp/server/ExtensionMcpProtocol.java | 14 + .../{HttpMcpProxy.java => HttpMcpClient.java} | 347 ++-- .../java/mcp/server/KnownProtocolVersion.java | 73 + .../smithy/java/mcp/server/McpCacheHint.java | 27 + .../java/mcp/server/McpCachePolicy.java | 57 + .../smithy/java/mcp/server/McpCacheScope.java | 27 + .../smithy/java/mcp/server/McpCall.java | 217 +++ .../smithy/java/mcp/server/McpCatalog.java | 904 ++++++++++ .../smithy/java/mcp/server/McpCursorPage.java | 14 + .../java/mcp/server/McpDomainOperations.java | 263 +++ .../smithy/java/mcp/server/McpEngine.java | 346 ++++ .../smithy/java/mcp/server/McpError.java | 20 + .../java/mcp/server/McpExecutionContext.java | 24 + .../java/mcp/server/McpExecutionHook.java | 60 - .../java/mcp/server/McpExtensionMethod.java | 30 + .../java/mcp/server/McpHttpBinding.java | 160 ++ .../java/mcp/server/McpHttpHandler.java | 269 +++ .../java/mcp/server/McpInterceptor.java | 77 + .../java/mcp/server/McpInterceptorChain.java | 101 ++ .../smithy/java/mcp/server/McpJson.java | 20 + .../smithy/java/mcp/server/McpMetadata.java | 65 + .../smithy/java/mcp/server/McpMethod.java | 104 ++ .../smithy/java/mcp/server/McpOperations.java | 43 + .../smithy/java/mcp/server/McpOutcome.java | 33 + .../smithy/java/mcp/server/McpPage.java | 39 + .../java/mcp/server/McpPromptDescriptor.java | 8 + .../smithy/java/mcp/server/McpProtocol.java | 155 ++ .../java/mcp/server/McpProtocolException.java | 36 + .../java/mcp/server/McpProtocolFeatures.java | 23 + .../smithy/java/mcp/server/McpProtocolId.java | 26 + .../java/mcp/server/McpProtocolProvider.java | 20 + .../java/mcp/server/McpProtocolRegistry.java | 205 +++ .../java/mcp/server/McpRemoteClient.java | 477 +++++ .../java/mcp/server/McpRemoteException.java | 19 + .../java/mcp/server/McpRequestContext.java | 25 + .../java/mcp/server/McpRequestDecoder.java | 187 ++ .../java/mcp/server/McpSchemaFactory.java | 363 ++++ .../smithy/java/mcp/server/McpServer.java | 167 -- .../java/mcp/server/McpServerBuilder.java | 127 -- .../java/mcp/server/McpServerIdentity.java | 20 + .../java/mcp/server/McpServerInterceptor.java | 217 --- .../mcp/server/McpServerInterceptorChain.java | 136 -- .../java/mcp/server/McpServerProxy.java | 240 --- .../smithy/java/mcp/server/McpService.java | 1588 ----------------- .../smithy/java/mcp/server/McpSession.java | 44 + .../java/mcp/server/McpSourceSnapshot.java | 13 + .../smithy/java/mcp/server/McpSources.java | 53 + .../java/mcp/server/McpToolCallHook.java | 67 - .../java/mcp/server/McpToolDescriptor.java | 26 + .../mcp/server/McpToolExecutionContext.java | 29 + .../java/mcp/server/McpToolExecutor.java | 168 ++ .../java/mcp/server/McpTransportContext.java | 38 + .../server/McpUnsupportedMethodException.java | 15 + .../smithy/java/mcp/server/McpWireCodec.java | 149 ++ .../smithy/java/mcp/server/McpWireNames.java | 15 + .../amazon/smithy/java/mcp/server/Prompt.java | 38 +- .../java/mcp/server/ProtocolVersion.java | 101 +- .../mcp/server/SmithyDocumentAdapter.java | 254 +++ .../{StdioProxy.java => StdioMcpClient.java} | 182 +- .../java/mcp/server/StdioMcpServer.java | 206 +++ .../mcp/server/StdioMcpServerBuilder.java | 166 ++ .../mcp/server/UnknownProtocolVersion.java | 22 + .../smithy/java/mcp/server/package-info.java | 13 +- ...pProxyTest.java => HttpMcpClientTest.java} | 528 +++++- .../java/mcp/server/McpArchitectureTest.java | 615 +++++++ .../java/mcp/server/McpCatalogTest.java | 780 ++++++++ .../java/mcp/server/McpHttpHandlerTest.java | 265 +++ .../java/mcp/server/McpRemoteClientTest.java | 341 ++++ .../java/mcp/server/McpSchemaFactoryTest.java | 147 ++ .../java/mcp/server/McpServerProxyTest.java | 245 --- .../java/mcp/server/McpServiceTest.java | 243 --- .../java/mcp/server/ProtocolVersionTest.java | 32 +- .../java/mcp/server/StdioMcpClientTest.java | 62 + ...erverTest.java => StdioMcpServerTest.java} | 753 +++++--- .../java/mcp/server/StdioProxyTest.java | 48 - .../java/mcp/server/TestInputStream.java | 69 - .../server/spi/ExternalProtocolApiTest.java | 52 + .../server/utils/TestJavaCodegenRunner.java | 28 +- .../java/mcp/server/TestInputStream.java | 0 .../java/mcp/server/TestOutputStream.java | 0 smithy-ai-traits/model/mcp.smithy | 12 + .../amazon/smithy/ai/McpHeaderValidator.java | 42 + ...e.amazon.smithy.model.validation.Validator | 1 + .../smithy/ai/McpHeaderValidatorTest.java | 45 + .../test/resources/mcp-header-invalid.smithy | 10 + .../test/resources/mcp-header-valid.smithy | 10 + 103 files changed, 10444 insertions(+), 4013 deletions(-) create mode 100644 mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/McpConformanceTest.java rename mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/{McpServerIntegrationTest.java => StdioMcpServerIntegrationTest.java} (96%) delete mode 100644 mcp/mcp-server/src/it/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java create mode 100644 mcp/mcp-server/src/it/resources/META-INF/smithy/conformance.smithy create mode 100644 mcp/mcp-server/src/it/resources/conformance-baseline-2025-11-25.yaml create mode 100644 mcp/mcp-server/src/it/resources/conformance-baseline-2026-07-28.yaml create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocol.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocols.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/ExtensionMcpProtocol.java rename mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/{HttpMcpProxy.java => HttpMcpClient.java} (53%) create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/KnownProtocolVersion.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheHint.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCachePolicy.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheScope.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCall.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCatalog.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCursorPage.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpDomainOperations.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpEngine.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpError.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionContext.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExecutionHook.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpExtensionMethod.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpBinding.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpHandler.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptor.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptorChain.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpJson.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMetadata.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMethod.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOperations.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOutcome.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPage.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPromptDescriptor.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocol.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolException.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolFeatures.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolId.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolProvider.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolRegistry.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteClient.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteException.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestContext.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestDecoder.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSchemaFactory.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServer.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerBuilder.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerIdentity.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptor.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerInterceptorChain.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSession.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSourceSnapshot.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSources.java delete mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolCallHook.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolDescriptor.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutionContext.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutor.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpTransportContext.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpUnsupportedMethodException.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireCodec.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireNames.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/SmithyDocumentAdapter.java rename mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/{StdioProxy.java => StdioMcpClient.java} (63%) create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServer.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServerBuilder.java create mode 100644 mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/UnknownProtocolVersion.java rename mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/{HttpMcpProxyTest.java => HttpMcpClientTest.java} (56%) create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpArchitectureTest.java create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpCatalogTest.java create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpHttpHandlerTest.java create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpRemoteClientTest.java create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpSchemaFactoryTest.java delete mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerProxyTest.java delete mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpClientTest.java rename mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/{McpServerTest.java => StdioMcpServerTest.java} (80%) delete mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioProxyTest.java delete mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/TestInputStream.java create mode 100644 mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/spi/ExternalProtocolApiTest.java rename mcp/mcp-server/src/{it => testFixtures}/java/software/amazon/smithy/java/mcp/server/TestInputStream.java (100%) rename mcp/mcp-server/src/{test => testFixtures}/java/software/amazon/smithy/java/mcp/server/TestOutputStream.java (100%) create mode 100644 smithy-ai-traits/model/mcp.smithy create mode 100644 smithy-ai-traits/src/main/java/software/amazon/smithy/ai/McpHeaderValidator.java create mode 100644 smithy-ai-traits/src/test/java/software/amazon/smithy/ai/McpHeaderValidatorTest.java create mode 100644 smithy-ai-traits/src/test/resources/mcp-header-invalid.smithy create mode 100644 smithy-ai-traits/src/test/resources/mcp-header-valid.smithy 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 0473261238..0f1e5ac95e 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 @@ -156,6 +157,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..3eac523d96 100644 --- a/mcp/mcp-server/README.md +++ b/mcp/mcp-server/README.md @@ -5,3 +5,224 @@ > 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); + } + + @Override + public int initializationPriority() { + return 100; + } +} +``` + +`initializationPriority` is considered only for stateful protocols that support +`initialize`. Higher priorities are preferred when a client requests an unknown +or stateless version. Equal highest priorities fail engine construction rather +than making fallback depend on registration order. + +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. + +## Remote pagination + +Remote tool and prompt listings are lazy. `McpRemoteClient.listTools()` and +`listPrompts()` return an `McpPage` containing the current items and an optional +continuation. Fetching the continuation performs exactly one additional upstream +request: + +```java +var page = remoteClient.listTools(); +process(page.items()); + +while (page.nextPage().isPresent()) { + page = page.nextPage().orElseThrow().fetch(); + process(page.items()); +} +``` + +This makes aggregation explicit for direct client users. When a remote client is +attached to an MCP engine, the engine translates each continuation into an opaque +downstream `nextCursor`. It does not eagerly drain the remote listing. Descriptors +from pages already requested are retained for tool and prompt dispatch, while a +fresh listing still starts from the cached first page. This favors availability +over immediate removal: if an upstream deletes an item from a later page, the +previously advertised descriptor can remain dispatchable until pagination reaches +and reconciles that page again. + +`ToolFilter` is enforced for both `tools/list` and `tools/call`; a tool hidden +from discovery cannot be invoked by name through the same engine. + +The engine forwards the active protocol to remote listings. Stateless requests +include their required `_meta` fields on every page, so modern-only upstream +servers do not depend on a legacy initialization handshake. If a remote does not +support stateless discovery, the engine initializes that remote independently +using the proxy server's identity and the highest-priority compatible stateful +protocol. + +## HTTP parameter headers + +Annotate a Smithy input member with `smithy.ai#mcpHeader` to mirror that value in +an `Mcp-Param-*` HTTP header. The trait must be retained as a Java runtime trait +when generating the service: + +```json +{ + "runtimeTraits": [ + "smithy.ai#mcpHeader" + ] +} +``` + +The MCP integration and conformance build include this setting. Projects invoking +Smithy Java code generation directly must add it to their codegen settings. + +## Cache hints + +Stateless cacheable results default to `ttlMs: 0` and `cacheScope: "private"`. +Configure a different default or override individual methods with an immutable +cache policy: + +```java +var cachePolicy = McpCachePolicy.builder() + .hint( + McpMethod.Standard.TOOLS_LIST, + new McpCacheHint(30_000, McpCacheScope.PUBLIC)) + .build(); + +var engine = McpEngine.builder() + .cachePolicy(cachePolicy) + .build(); +``` + +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 a11025a72c..fa5d37560f 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) { @@ -134,34 +134,35 @@ private ToolSchemas getMcpEchoToolSchemas() { write("tools/list", Document.of(Map.of())); var responseJson = readRawResponse(); var toolsNode = OBJECT_MAPPER.readTree(responseJson).path("result").path("tools"); - for (var toolNode : toolsNode) { - if (toolNode.path("name").asString().equals("McpEcho")) { - var inputSchemaNode = toolNode.path("inputSchema"); - var outputSchemaNode = toolNode.path("outputSchema"); - return new ToolSchemas( - SCHEMA_FACTORY.getSchema(inputSchemaNode), - SCHEMA_FACTORY.getSchema(outputSchemaNode), - toolNode); - } - } - throw new AssertionError("McpEcho tool not found"); + var toolNode = findToolNode(toolsNode, "McpEcho"); + var inputSchemaNode = toolNode.path("inputSchema"); + var outputSchemaNode = toolNode.path("outputSchema"); + return new ToolSchemas( + SCHEMA_FACTORY.getSchema(inputSchemaNode), + SCHEMA_FACTORY.getSchema(outputSchemaNode), + toolNode); } 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() { @@ -188,22 +189,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(); @@ -214,7 +220,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(); @@ -428,12 +434,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 ========== @@ -832,7 +838,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) @@ -1293,8 +1299,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)); @@ -1319,8 +1325,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()); @@ -1448,7 +1454,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); @@ -1474,7 +1480,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")); @@ -2035,6 +2041,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(); @@ -2328,7 +2376,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); @@ -2361,7 +2409,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 @@ -2377,7 +2425,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..9b38ae5766 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/BuiltInProtocol.java @@ -0,0 +1,78 @@ +/* + * 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; +import software.amazon.smithy.model.shapes.ShapeType; + +/** + * 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 int initializationPriority() { + return version.ordinal(); + } + + @Override + public Document decorateResult( + Document result, + McpMethod method, + McpServerIdentity serverIdentity, + McpCachePolicy cachePolicy + ) { + if (!features.statelessResults()) { + return result; + } + + var members = new HashMap<>(result.asStringMap()); + members.put("resultType", Document.of("complete")); + + var existingMeta = members.get("_meta"); + var meta = existingMeta != null + && (existingMeta.isType(ShapeType.MAP) || existingMeta.isType(ShapeType.STRUCTURE)) + ? new HashMap<>(existingMeta.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)); + + var cacheableMethod = switch (method) { + case McpMethod.Standard.SERVER_DISCOVER, + McpMethod.Standard.TOOLS_LIST, + McpMethod.Standard.PROMPTS_LIST -> + true; + default -> false; + }; + if (cacheableMethod) { + var hint = cachePolicy.hint((McpMethod.Standard) method); + members.put("ttlMs", Document.of(hint.ttlMs())); + members.put("cacheScope", Document.of(hint.scope().wireValue())); + } + 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 53% 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..a1078f6132 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,21 @@ 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.model.shapes.ShapeType; 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 +50,11 @@ public final class HttpMcpProxy extends McpServerProxy { private final IdentityResolver identityResolver; private final Context signerContext; private final Duration timeout; - private volatile String sessionId; + private final AtomicReference>> toolHeaderParameters = + new AtomicReference<>(Map.of()); + private final AtomicReference sessionId = new AtomicReference<>(); - 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 +122,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 +138,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 +147,54 @@ public static Builder builder() { } @Override - public CompletableFuture rpc(JsonRpcRequest request) { + protected void onToolsPage(List tools) { + toolHeaderParameters.updateAndGet(current -> { + var updatedMappings = new HashMap<>(current); + for (var tool : tools) { + var mappings = McpHttpBinding.headerParameters(tool); + if (mappings.isEmpty()) { + updatedMappings.remove(tool.getName()); + } else { + updatedMappings.put(tool.getName(), mappings); + } + } + return Map.copyOf(updatedMappings); + }); + } + + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + return exchange(request, true); + } + + private JsonRpcResponse exchange(JsonRpcRequest request, boolean allowSessionRecovery) { 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; + String currentSessionId = sessionId.get(); + long currentGeneration = initializationGeneration(); 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,36 +219,115 @@ 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; + sessionId.set(responseSessionId); LOG.debug("Stored session ID from initialize response: {}", responseSessionId); } } // "When a client receives HTTP 404 in response to a request containing an Mcp-Session-Id, // it MUST start a new session by sending a new InitializeRequest without a session ID attached." - if (response.statusCode() == 404 && sessionId != null) { + if (response.statusCode() == 404 && currentSessionId != null) { LOG.debug("Received 404 with active session ID. Clearing session to force restart."); - sessionId = null; + var ownsRecovery = sessionId.compareAndSet(currentSessionId, null); + if (allowSessionRecovery + && !McpHttpBinding.isInitialize(request)) { + ByteBufferUtils.getBytes(response.body().asByteBuffer()); + if (restartSession(currentGeneration)) { + return exchange(request, false); + } + return sessionRecoveryFailure(request, ownsRecovery); + } } 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 JsonRpcResponse sessionRecoveryFailure( + JsonRpcRequest request, + boolean ownedRecovery + ) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .error(JsonRpcErrorResponse.builder() + .code(UPSTREAM_HTTP_ERROR_CODE) + .message(ownedRecovery + ? "HTTP 404: MCP session expired and could not be reinitialized" + : "HTTP 404: MCP session recovery did not complete") + .build()) + .build(); + } + + private McpProtocol requestProtocol(JsonRpcRequest request) { + var params = request.getParams(); + var meta = params == null + || !(params.isType(ShapeType.MAP) || params.isType(ShapeType.STRUCTURE)) + ? null + : params.getMember("_meta"); + var requestedVersion = meta == null + || !(meta.isType(ShapeType.MAP) || meta.isType(ShapeType.STRUCTURE)) + ? 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 +346,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); + 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, request.getId()); + 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'); } } - - // 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); - } - } + var result = processSseEvent(data, request.getId()); + if (result != null) { + return result; } - - 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(); - } - - return finalResponse; } catch (Exception e) { LOG.error("Error parsing SSE response", e); return JsonRpcResponse.builder() @@ -338,27 +376,78 @@ 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, Document expectedId) { + 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; + } + var response = document.asShape(JsonRpcResponse.builder()); + if (!requestIdsMatch(expectedId, response.getId())) { + LOG.warn("Ignoring SSE response with unexpected request ID"); + return null; + } + return response; + } catch (RuntimeException e) { + LOG.warn("Failed to parse SSE message: {}", jsonData, e); + return null; + } + } + + private boolean requestIdsMatch(Document expected, Document actual) { + if (expected == null || actual == null) { + return expected == actual; + } + try { + return StdioMcpClient.requestKey(expected).equals(StdioMcpClient.requestKey(actual)); + } catch (IllegalStateException e) { + return false; + } } - 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 +461,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 +476,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/McpCacheHint.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheHint.java new file mode 100644 index 0000000000..cfa404ac7c --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheHint.java @@ -0,0 +1,27 @@ +/* + * 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; + +/** + * Cache metadata attached to a cacheable MCP result. + * + * @param ttlMs cache lifetime in milliseconds. + * @param scope cache visibility. + */ +@SmithyUnstableApi +public record McpCacheHint(long ttlMs, McpCacheScope scope) { + public static final McpCacheHint NO_CACHE = new McpCacheHint(0, McpCacheScope.PRIVATE); + + public McpCacheHint { + if (ttlMs < 0) { + throw new IllegalArgumentException("MCP cache TTL must not be negative"); + } + Objects.requireNonNull(scope, "scope"); + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCachePolicy.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCachePolicy.java new file mode 100644 index 0000000000..b8a88fcff7 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCachePolicy.java @@ -0,0 +1,57 @@ +/* + * 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.EnumMap; +import java.util.Map; +import java.util.Objects; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * Immutable cache-hint policy for MCP methods. + */ +@SmithyUnstableApi +public final class McpCachePolicy { + public static final McpCachePolicy DEFAULT = builder().build(); + + private final McpCacheHint defaultHint; + private final Map methodHints; + + private McpCachePolicy(Builder builder) { + defaultHint = builder.defaultHint; + methodHints = Map.copyOf(builder.methodHints); + } + + public McpCacheHint hint(McpMethod.Standard method) { + return methodHints.getOrDefault(method, defaultHint); + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private final Map methodHints = + new EnumMap<>(McpMethod.Standard.class); + private McpCacheHint defaultHint = McpCacheHint.NO_CACHE; + + public Builder defaultHint(McpCacheHint hint) { + defaultHint = Objects.requireNonNull(hint, "hint"); + return this; + } + + public Builder hint(McpMethod.Standard method, McpCacheHint hint) { + methodHints.put( + Objects.requireNonNull(method, "method"), + Objects.requireNonNull(hint, "hint")); + return this; + } + + public McpCachePolicy build() { + return new McpCachePolicy(this); + } + } +} diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheScope.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheScope.java new file mode 100644 index 0000000000..b38f5ecdd5 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCacheScope.java @@ -0,0 +1,27 @@ +/* + * 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; + +/** + * Visibility of a cached MCP result. + */ +@SmithyUnstableApi +public enum McpCacheScope { + PRIVATE("private"), + PUBLIC("public"); + + private final String wireValue; + + McpCacheScope(String wireValue) { + this.wireValue = wireValue; + } + + public String wireValue() { + return wireValue; + } +} 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..c1ba6fbc0e --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCatalog.java @@ -0,0 +1,904 @@ +/* + * 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.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.LongSupplier; +import java.util.function.UnaryOperator; +import software.amazon.smithy.java.core.schema.SchemaIndex; +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.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 static final int MAX_ACTIVE_CURSORS = 1_024; + private static final long CURSOR_TTL_NANOS = TimeUnit.MINUTES.toNanos(5); + private static final long REMOTE_RETRY_COOLDOWN_NANOS = TimeUnit.SECONDS.toNanos(30); + + private final AtomicReference state; + private final McpProtocolRegistry protocols; + private final McpServerIdentity identity; + private final LongSupplier nanoTime; + private final long remoteRetryCooldownNanos; + private final AtomicReference> remoteStart = new AtomicReference<>(); + private final Map remoteCatalogLoads = + new ConcurrentHashMap<>(); + private final ExecutorService notificationRefreshes = Executors.newVirtualThreadPerTaskExecutor(); + private final Map toolRefreshStates = new ConcurrentHashMap<>(); + private final Map promptRefreshStates = new ConcurrentHashMap<>(); + private final Map> toolCursors = new ConcurrentHashMap<>(); + private final Map> promptCursors = new ConcurrentHashMap<>(); + private final CopyOnWriteArrayList> notificationWriters = + new CopyOnWriteArrayList<>(); + private final CopyOnWriteArrayList> responseWriters = + new CopyOnWriteArrayList<>(); + private volatile JsonRpcRequest initializeRequest; + private volatile McpProtocol initializeProtocol; + + McpCatalog(Map services, List remoteClients) { + this( + services, + remoteClients, + McpProtocolRegistry.create(List.of(), List.of(), false), + new McpServerIdentity("mcp-server", "1.0.0"), + System::nanoTime, + REMOTE_RETRY_COOLDOWN_NANOS); + } + + McpCatalog( + Map services, + List remoteClients, + McpProtocolRegistry protocols, + McpServerIdentity identity + ) { + this( + services, + remoteClients, + protocols, + identity, + System::nanoTime, + REMOTE_RETRY_COOLDOWN_NANOS); + } + + McpCatalog( + Map services, + List remoteClients, + McpProtocolRegistry protocols, + McpServerIdentity identity, + LongSupplier nanoTime, + long remoteRetryCooldownNanos + ) { + this.protocols = protocols; + this.identity = identity; + this.nanoTime = nanoTime; + this.remoteRetryCooldownNanos = remoteRetryCooldownNanos; + 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); + var localSnapshot = createLocalSnapshot(immutableServices, Map.of(), Map.of()); + state = new AtomicReference<>(new CatalogState( + immutableServices, + clients, + Map.of(), + Map.of(), + localSnapshot.tools(), + localSnapshot.prompts(), + localSnapshot)); + } + + @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 McpCursorPage listTools(String cursor) { + if (cursor == null) { + var current = state.get(); + var pending = current.toolContinuations() + .entrySet() + .stream() + .map(entry -> new PendingPage<>(entry.getKey(), entry.getValue())) + .toList(); + return new McpCursorPage<>( + List.copyOf(current.initialTools().values()), + registerCursor(toolCursors, pending)); + } + + var cursorState = requireCursor(toolCursors, cursor, "tools"); + var pending = new ArrayList<>(cursorState.pending()); + var current = pending.removeFirst(); + var page = current.nextPage().fetch(); + toolCursors.remove(cursor, cursorState); + page.nextPage().ifPresent(next -> pending.add(new PendingPage<>(current.client(), next))); + var descriptors = mergeRemoteToolPage(current.client(), page); + return new McpCursorPage<>(descriptors, registerCursor(toolCursors, pending)); + } + + @Override + public McpCursorPage listPrompts(String cursor) { + if (cursor == null) { + var current = state.get(); + var pending = current.promptContinuations() + .entrySet() + .stream() + .map(entry -> new PendingPage<>(entry.getKey(), entry.getValue())) + .toList(); + return new McpCursorPage<>( + List.copyOf(current.initialPrompts().values()), + registerCursor(promptCursors, pending)); + } + + var cursorState = requireCursor(promptCursors, cursor, "prompts"); + var pending = new ArrayList<>(cursorState.pending()); + var current = pending.removeFirst(); + var page = current.nextPage().fetch(); + promptCursors.remove(cursor, cursorState); + page.nextPage().ifPresent(next -> pending.add(new PendingPage<>(current.client(), next))); + var descriptors = mergeRemotePromptPage(current.client(), page); + return new McpCursorPage<>(descriptors, registerCursor(promptCursors, pending)); + } + + @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 + ) { + notificationWriters.addIfAbsent(notificationWriter); + responseWriters.addIfAbsent(responseWriter); + runOnce(remoteStart, () -> forEachRemoteInParallel("start", McpRemoteClient::start)); + } + + @Override + public void initializeRemoteClients(McpProtocol protocol) { + var request = initializeRequest(protocol); + initializeRequest = request; + initializeProtocol = protocol; + forEachRemoteInParallel( + "initialize", + client -> loadRemoteCatalog(client, protocol)); + } + + @Override + public void ensureRemoteCatalogLoaded(McpProtocol requestedProtocol) { + var protocol = catalogProtocol(requestedProtocol); + forEachRemoteInParallel( + "refresh", + client -> loadRemoteCatalog(client, protocol)); + } + + private McpProtocol catalogProtocol(McpProtocol requestedProtocol) { + var initialized = initializeProtocol; + return initialized != null && !requestedProtocol.usesStatelessMetadata() + ? initialized + : requestedProtocol; + } + + @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))); + + var initialTools = new HashMap<>(remoteTools(current.initialTools())); + initialTools.putAll(schemaFactory.createTools(services)); + var initialPrompts = createPromptSnapshot( + services, + remotePrompts(current.initialPrompts())); + + return new CatalogState( + services, + current.remoteClients(), + current.toolContinuations(), + current.promptContinuations(), + initialTools, + initialPrompts, + 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.toolContinuations(), + current.promptContinuations(), + current.initialTools(), + current.initialPrompts(), + current.snapshot()); + }); + + try { + client.start(); + var currentInitializeRequest = initializeRequest; + if (currentInitializeRequest != null) { + loadRemoteCatalog(client, initializeProtocol); + } else { + loadRemoteCatalog(client, protocols.defaultProtocol()); + } + } 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 initializeRemote( + McpRemoteClient client, + JsonRpcRequest request, + McpProtocol protocol + ) { + client.initialize( + this::writeResponse, + notification -> onRemoteNotification(client, notification), + request, + protocol, + protocols); + } + + private void loadRemoteCatalog(McpRemoteClient client, McpProtocol requestedProtocol) { + var key = new RemoteCatalogKey(client, requestedProtocol.id()); + var load = remoteCatalogLoads.computeIfAbsent(key, ignored -> new RemoteLoadState()); + if (load.retryAfterNanos().get() > nanoTime.getAsLong()) { + return; + } + + runOnce(load.operation(), () -> { + try { + var protocol = prepareRemoteClient(client, requestedProtocol); + client.usingProtocol(protocol, () -> { + refresh(client); + return null; + }); + load.retryAfterNanos().set(0); + } catch (RuntimeException e) { + load.retryAfterNanos().set(nanoTime.getAsLong() + remoteRetryCooldownNanos); + throw e; + } + }); + } + + private McpProtocol prepareRemoteClient( + McpRemoteClient client, + McpProtocol requestedProtocol + ) { + if (requestedProtocol.usesStatelessMetadata()) { + if (client.supportsStatelessProtocol(requestedProtocol)) { + return requestedProtocol; + } + if (client.initialized()) { + return client.negotiatedProtocol(); + } + } else if (client.initialized()) { + return client.negotiatedProtocol(); + } + + var initializationProtocol = requestedProtocol.supportedMethods() + .contains(McpMethod.Standard.INITIALIZE) + ? requestedProtocol + : protocols.initializationFallbackProtocol(); + var request = initializeRequest(initializationProtocol); + initializeRemote(client, request, initializationProtocol); + return client.negotiatedProtocol(); + } + + private JsonRpcRequest initializeRequest(McpProtocol protocol) { + return JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(0)) + .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( + "name", + Document.of(identity.name()), + "version", + Document.of(identity.version())))))) + .build(); + } + + private void onRemoteNotification(McpRemoteClient client, JsonRpcRequest notification) { + if (McpMethod.Standard.NOTIFICATIONS_TOOLS_LIST_CHANGED.wireName().equals(notification.getMethod())) { + scheduleRefresh(client, toolRefreshStates, this::refreshTools); + } else if (McpMethod.Standard.NOTIFICATIONS_PROMPTS_LIST_CHANGED.wireName().equals(notification.getMethod())) { + scheduleRefresh(client, promptRefreshStates, this::refreshPrompts); + } + notificationWriters.forEach(writer -> writer.accept(notification)); + } + + private void writeResponse(JsonRpcResponse response) { + responseWriters.forEach(writer -> writer.accept(response)); + } + + private void scheduleRefresh( + McpRemoteClient client, + Map states, + Consumer refresh + ) { + var state = states.computeIfAbsent(client, ignored -> new RefreshState()); + if (state.request()) { + notificationRefreshes.submit(() -> runScheduledRefreshes(client, state, refresh)); + } + } + + private void runScheduledRefreshes( + McpRemoteClient client, + RefreshState state, + Consumer refresh + ) { + while (state.takeRequest()) { + refresh.accept(client); + } + } + + private void refresh(McpRemoteClient client) { + McpPage remoteTools = McpPage.last(List.of()); + boolean toolsLoaded = false; + try { + remoteTools = client.listTools(); + toolsLoaded = true; + } catch (RuntimeException e) { + LOG.error("Failed to refresh tools from remote MCP client: " + client.name(), e); + } + + McpPage remotePrompts = McpPage.last(List.of()); + boolean promptsLoaded = false; + try { + remotePrompts = 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); + if (!toolsLoaded && !promptsLoaded) { + throw new McpRemoteException( + "Remote MCP client did not provide a tools or prompts catalog: " + client.name()); + } + } + + private void refreshTools(McpRemoteClient client) { + try { + mergeRemoteSnapshot( + client, + client.listTools(), + true, + McpPage.last(List.of()), + false); + } catch (RuntimeException e) { + LOG.error("Failed to refresh tools from remote MCP client: " + client.name(), e); + } + } + + private void refreshPrompts(McpRemoteClient client) { + try { + mergeRemoteSnapshot( + client, + McpPage.last(List.of()), + false, + client.listPrompts(), + true); + } catch (RuntimeException e) { + LOG.error("Failed to refresh prompts from remote MCP client: " + client.name(), e); + } + } + + private void mergeRemoteSnapshot( + McpRemoteClient client, + McpPage remoteTools, + boolean toolsLoaded, + McpPage remotePrompts, + boolean promptsLoaded + ) { + updateState(current -> { + var tools = new HashMap<>(current.snapshot().tools()); + var initialTools = new HashMap<>(current.initialTools()); + var toolContinuations = new HashMap<>(current.toolContinuations()); + if (toolsLoaded) { + initialTools.forEach((name, descriptor) -> { + if (descriptor.target() instanceof McpToolDescriptor.RemoteTarget remote + && remote.client() == client) { + tools.computeIfPresent(name, + ( + ignored, + existing) -> existing + .target() instanceof McpToolDescriptor.RemoteTarget existingRemote + && existingRemote.client() == client + ? null + : existing); + } + }); + for (var info : remoteTools.items()) { + putRemoteTool(tools, + new McpToolDescriptor( + info, + client.name(), + new McpToolDescriptor.RemoteTarget(client), + McpHttpBinding.headerParameters(info))); + } + initialTools.entrySet() + .removeIf(entry -> entry.getValue().target() instanceof McpToolDescriptor.RemoteTarget remote + && remote.client() == client); + for (var info : remoteTools.items()) { + putRemoteTool(initialTools, + new McpToolDescriptor( + info, + client.name(), + new McpToolDescriptor.RemoteTarget(client), + McpHttpBinding.headerParameters(info))); + } + replaceContinuation(toolContinuations, client, remoteTools.nextPage()); + } + + var prompts = new HashMap<>(current.snapshot().prompts()); + var initialPrompts = new HashMap<>(current.initialPrompts()); + var promptContinuations = new HashMap<>(current.promptContinuations()); + if (promptsLoaded) { + initialPrompts.forEach((name, descriptor) -> { + if (descriptor.remoteClient() == client) { + prompts.computeIfPresent(name, + (ignored, existing) -> existing.remoteClient() == client ? null : existing); + } + }); + for (var info : remotePrompts.items()) { + putRemotePrompt( + prompts, + PromptLoader.normalize(info.getName()), + new McpPromptDescriptor(new Prompt(info, client), client)); + } + initialPrompts.entrySet().removeIf(entry -> entry.getValue().remoteClient() == client); + for (var info : remotePrompts.items()) { + putRemotePrompt( + initialPrompts, + PromptLoader.normalize(info.getName()), + new McpPromptDescriptor(new Prompt(info, client), client)); + } + replaceContinuation(promptContinuations, client, remotePrompts.nextPage()); + } + + return new CatalogState( + current.services(), + current.remoteClients(), + toolContinuations, + promptContinuations, + initialTools, + initialPrompts, + new McpSourceSnapshot( + Map.copyOf(tools), + Map.copyOf(prompts), + current.snapshot().documentAdapter())); + }); + } + + private List mergeRemoteToolPage( + McpRemoteClient client, + McpPage page + ) { + var descriptors = page.items() + .stream() + .map(info -> new McpToolDescriptor( + info, + client.name(), + new McpToolDescriptor.RemoteTarget(client), + McpHttpBinding.headerParameters(info))) + .toList(); + updateState(current -> { + var tools = new HashMap<>(current.snapshot().tools()); + descriptors.forEach(descriptor -> putRemoteTool(tools, descriptor)); + return new CatalogState( + current.services(), + current.remoteClients(), + current.toolContinuations(), + current.promptContinuations(), + current.initialTools(), + current.initialPrompts(), + new McpSourceSnapshot( + Map.copyOf(tools), + current.snapshot().prompts(), + current.snapshot().documentAdapter())); + }); + return descriptors; + } + + private List mergeRemotePromptPage( + McpRemoteClient client, + McpPage page + ) { + var descriptors = page.items() + .stream() + .map(info -> new McpPromptDescriptor(new Prompt(info, client), client)) + .toList(); + updateState(current -> { + var prompts = new HashMap<>(current.snapshot().prompts()); + descriptors.forEach(descriptor -> putRemotePrompt( + prompts, + PromptLoader.normalize(descriptor.prompt().promptInfo().getName()), + descriptor)); + return new CatalogState( + current.services(), + current.remoteClients(), + current.toolContinuations(), + current.promptContinuations(), + current.initialTools(), + current.initialPrompts(), + new McpSourceSnapshot( + current.snapshot().tools(), + Map.copyOf(prompts), + current.snapshot().documentAdapter())); + }); + return descriptors; + } + + private void replaceContinuation( + Map> continuations, + McpRemoteClient client, + Optional> nextPage + ) { + if (nextPage.isPresent()) { + continuations.put(client, nextPage.get()); + } else { + continuations.remove(client); + } + } + + private void putRemoteTool( + Map tools, + McpToolDescriptor candidate + ) { + var name = candidate.info().getName(); + var existing = tools.putIfAbsent(name, candidate); + if (existing != null + && existing.target() instanceof McpToolDescriptor.RemoteTarget existingRemote + && candidate.target() instanceof McpToolDescriptor.RemoteTarget candidateRemote + && existingRemote.client() == candidateRemote.client()) { + tools.replace(name, existing, candidate); + return; + } + if (existing != null) { + LOG.warn( + "Ignoring remote MCP tool {} from {} because it is already provided by {}", + name, + candidate.serverId(), + existing.serverId()); + } + } + + private void putRemotePrompt( + Map prompts, + String name, + McpPromptDescriptor candidate + ) { + var existing = prompts.putIfAbsent(name, candidate); + if (existing != null && existing.remoteClient() == candidate.remoteClient()) { + prompts.replace(name, existing, candidate); + } + } + + private PageCursor requireCursor( + Map> cursors, + String cursor, + String listing + ) { + var state = cursors.get(cursor); + if (state == null || state.expired(System.nanoTime())) { + if (state != null) { + cursors.remove(cursor, state); + } + throw new McpProtocolException(-32602, "Invalid or expired " + listing + " cursor"); + } + return state; + } + + private String registerCursor( + Map> cursors, + List> pending + ) { + if (pending.isEmpty()) { + return null; + } + evictExpiredAndOverflow(cursors); + var cursor = UUID.randomUUID().toString(); + cursors.put(cursor, new PageCursor<>(pending, System.nanoTime())); + return cursor; + } + + private void evictExpiredAndOverflow(Map> cursors) { + var now = System.nanoTime(); + cursors.entrySet().removeIf(entry -> entry.getValue().expired(now)); + while (cursors.size() >= MAX_ACTIVE_CURSORS) { + var oldest = cursors.entrySet() + .stream() + .min(Map.Entry.comparingByValue( + (left, right) -> Long.compare(left.createdAtNanos(), right.createdAtNanos()))) + .orElse(null); + if (oldest == null || !cursors.remove(oldest.getKey(), oldest.getValue())) { + return; + } + } + } + + 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<>(remoteTools); + tools.putAll(schemaFactory.createTools(services)); + + 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, + Map> toolContinuations, + Map> promptContinuations, + Map initialTools, + Map initialPrompts, + McpSourceSnapshot snapshot) { + private CatalogState { + services = Map.copyOf(services); + remoteClients = Map.copyOf(remoteClients); + toolContinuations = Map.copyOf(toolContinuations); + promptContinuations = Map.copyOf(promptContinuations); + initialTools = Map.copyOf(initialTools); + initialPrompts = Map.copyOf(initialPrompts); + } + } + + private record PendingPage(McpRemoteClient client, McpPage.NextPage nextPage) {} + + private record RemoteCatalogKey(McpRemoteClient client, McpProtocolId protocol) {} + + private record RemoteLoadState( + AtomicReference> operation, + AtomicLong retryAfterNanos) { + private RemoteLoadState() { + this(new AtomicReference<>(), new AtomicLong()); + } + } + + private record PageCursor(List> pending, long createdAtNanos) { + private PageCursor { + pending = List.copyOf(pending); + } + + boolean expired(long now) { + return now - createdAtNanos >= CURSOR_TTL_NANOS; + } + } + + 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/McpCursorPage.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCursorPage.java new file mode 100644 index 0000000000..44adde9df3 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpCursorPage.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 java.util.List; + +record McpCursorPage(List items, String nextCursor) { + McpCursorPage { + items = List.copyOf(items); + } +} 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..798f159097 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpDomainOperations.java @@ -0,0 +1,263 @@ +/* + * 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 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; +import software.amazon.smithy.model.shapes.ShapeType; + +/** + * 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, + toolFilter); + } + + @Override + public McpOutcome initialize(McpCall.Initialize call, McpRequestContext context) { + observeInitialize(call); + var protocol = protocols.require(context.protocolVersion()); + sources.initializeRemoteClients(protocol); + + var result = InitializeResult.builder() + .protocolVersion(context.protocolVersion().identifier()) + .capabilities(initializeCapabilities(protocol, context.transport())) + .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) { + var protocol = protocols.require(context.protocolVersion()); + sources.ensureRemoteCatalogLoaded(protocol); + var capabilities = discoverCapabilities(protocol); + 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) { + var protocol = protocols.require(context.protocolVersion()); + sources.ensureRemoteCatalogLoaded(protocol); + var page = sources.listTools(call.cursor()); + var tools = page.items() + .stream() + .filter(tool -> toolFilter.allowTool(tool.serverId(), tool.info().getName())) + .map(tool -> projectTool(tool.info(), protocol)) + .toList(); + var result = ListToolsResult.builder().tools(tools); + if (page.nextCursor() != null) { + result.nextCursor(page.nextCursor()); + } + return new McpOutcome.Success( + call.id(), + Document.of(result.build())); + } + + @Override + public McpOutcome callTool(McpCall.CallTool call, McpRequestContext context) { + sources.ensureRemoteCatalogLoaded(protocols.require(context.protocolVersion())); + 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(protocols.require(context.protocolVersion())); + var page = sources.listPrompts(call.cursor()); + var prompts = page.items() + .stream() + .map(descriptor -> descriptor.prompt().promptInfo()) + .toList(); + var result = ListPromptsResult.builder().prompts(prompts); + if (page.nextCursor() != null) { + result.nextCursor(page.nextCursor()); + } + return new McpOutcome.Success( + call.id(), + Document.of(result.build())); + } + + @Override + public McpOutcome getPrompt(McpCall.GetPrompt call, McpRequestContext context) { + var protocol = protocols.require(context.protocolVersion()); + sources.ensureRemoteCatalogLoaded(protocol); + 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(), + call.metadata(), + protocol))); + } + + private Capabilities initializeCapabilities( + McpProtocol protocol, + McpTransportContext transport + ) { + var builder = Capabilities.builder(); + if (supports(protocol, McpMethod.Standard.TOOLS_LIST)) { + var tools = Tools.builder(); + if (transport.supportsServerNotifications()) { + tools.listChanged(true); + } + builder.tools(tools.build()); + } + if (supports(protocol, McpMethod.Standard.PROMPTS_LIST)) { + var prompts = Prompts.builder(); + if (transport.supportsServerNotifications()) { + prompts.listChanged(true); + } + builder.prompts(prompts.build()); + } + return builder.build(); + } + + private Document discoverCapabilities(McpProtocol protocol) { + var capabilities = new HashMap(); + if (supports(protocol, McpMethod.Standard.TOOLS_LIST)) { + capabilities.put("tools", Document.of(Map.of())); + } + if (supports(protocol, McpMethod.Standard.PROMPTS_LIST)) { + capabilities.put("prompts", Document.of(Map.of())); + } + return Document.of(capabilities); + } + + private boolean supports(McpProtocol protocol, McpMethod.Standard method) { + return protocol.supportedMethods().contains(method); + } + + @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())); + } + + 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(); + var roots = objectMember(capabilities, "roots"); + var listChanged = objectMember(roots, "listChanged"); + boolean rootsListChanged = listChanged != null + && listChanged.isType(ShapeType.BOOLEAN) + && listChanged.asBoolean(); + boolean sampling = objectMember(capabilities, "sampling") != null; + boolean elicitation = objectMember(capabilities, "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) { + var member = objectMember(document, name); + return member != null && member.isType(ShapeType.STRING) ? member.asString() : null; + } + + private Document objectMember(Document document, String name) { + return McpHttpBinding.isObject(document) ? document.getMember(name) : null; + } +} 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..22d978a6a9 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpEngine.java @@ -0,0 +1,346 @@ +/* + * 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 final McpCachePolicy cachePolicy; + + 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; + cachePolicy = builder.cachePolicy; + sources = new McpCatalog(builder.services, builder.remoteClients, protocols, identity); + 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, cachePolicy)); + } + } 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) { + if (request.getId() == null) { + return McpOutcome.NoResponse.INSTANCE; + } + 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; + private McpCachePolicy cachePolicy = McpCachePolicy.DEFAULT; + + 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 Builder cachePolicy(McpCachePolicy cachePolicy) { + this.cachePolicy = Objects.requireNonNull(cachePolicy, "cachePolicy"); + 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..47c9f56411 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpBinding.java @@ -0,0 +1,160 @@ +/* + * 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 (!isObject(params)) { + 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 = isObject(params) ? params.getMember("_meta") : null; + return isObject(metadata) ? stringMember(metadata, McpWireNames.PROTOCOL_VERSION) : null; + } + + static String stringMember(Document document, String name) { + if (!isObject(document)) { + return null; + } + var member = document.getMember(name); + return member == null || !member.isType(ShapeType.STRING) ? null : member.asString(); + } + + static boolean isObject(Document document) { + return document != null + && (document.isType(ShapeType.MAP) || document.isType(ShapeType.STRUCTURE)); + } + + 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 (response.getError().getCode() == -32700 || response.getError().getCode() == -32600) { + 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..1fe0b991d7 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpHttpHandler.java @@ -0,0 +1,269 @@ +/* + * 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 java.util.concurrent.atomic.AtomicBoolean; +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; + private final AtomicBoolean bound = new AtomicBoolean(); + + public McpHttpHandler(McpEngine engine) { + this(engine, false); + } + + private McpHttpHandler(McpEngine engine, boolean loopbackOnly) { + this.engine = engine; + this.loopbackOnly = loopbackOnly; + } + + /** + * 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 + ) { + if (bound.compareAndSet(false, true)) { + engine.bindTransport(ignored -> {}, ignored -> {}); + } + 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 identifier = 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 = McpHttpBinding.isObject(arguments) + ? arguments.getMember(parameterName) + : null; + 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..3e887ba801 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpInterceptorChain.java @@ -0,0 +1,101 @@ +/* + * 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.Objects; + +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) { + var modified = interceptor.modifyAfterExecution(context, current, currentError); + if (modified != null) { + current = modified; + currentError = null; + } + } + if (currentError != null) { + throw currentError; + } + return Objects.requireNonNull(current, "MCP interceptor returned a null outcome"); + } + + @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) { + var modified = interceptor.modifyAfterToolCall(context, current, currentError); + if (modified != null) { + current = modified; + currentError = null; + } + } + if (currentError != null) { + throw currentError; + } + return Objects.requireNonNull(current, "MCP interceptor returned a null tool outcome"); + } +} 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..94e5382b43 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMetadata.java @@ -0,0 +1,65 @@ +/* + * 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.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); + } + + static McpMetadata forProtocol(McpProtocol protocol) { + return protocol.usesStatelessMetadata() + ? new McpMetadata( + protocol.protocolVersion(), + null, + Document.of(Map.of()), + Map.of()) + : EMPTY; + } + + Document applyTo(Document params) { + if (this == EMPTY + || (protocolVersion == null + && clientInfo == null + && clientCapabilities == null + && extensions.isEmpty())) { + return params; + } + + var values = params == null + ? new HashMap() + : new HashMap<>(params.asStringMap()); + var meta = new HashMap<>(extensions); + if (protocolVersion != null) { + meta.put(McpWireNames.PROTOCOL_VERSION, Document.of(protocolVersion.identifier())); + } + if (clientInfo != null) { + meta.put(McpWireNames.CLIENT_INFO, clientInfo); + } + if (clientCapabilities != null) { + meta.put(McpWireNames.CLIENT_CAPABILITIES, 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/McpMethod.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMethod.java new file mode 100644 index 0000000000..d7c74d89f3 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpMethod.java @@ -0,0 +1,104 @@ +/* + * 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_PROMPTS_LIST_CHANGED("notifications/prompts/list_changed"), + 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..60c13aec1f --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpOperations.java @@ -0,0 +1,43 @@ +/* + * 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); + + default McpOutcome complete(McpCall.Complete call, McpRequestContext context) { + throw new UnsupportedOperationException("completion/complete is not implemented"); + } + + default McpOutcome setLogLevel(McpCall.SetLogLevel call, McpRequestContext context) { + throw new UnsupportedOperationException("logging/setLevel is not implemented"); + } + + default McpOutcome readResource(McpCall.ReadResource call, McpRequestContext context) { + throw new UnsupportedOperationException("resources/read is not implemented"); + } +} 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/McpPage.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPage.java new file mode 100644 index 0000000000..f47b04e4a9 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpPage.java @@ -0,0 +1,39 @@ +/* + * 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.Optional; +import software.amazon.smithy.utils.SmithyUnstableApi; + +/** + * One page returned by a remote MCP listing operation. + * + *

The continuation is deliberately executable rather than exposing the remote + * server's opaque cursor. Calling it performs exactly one request for the next page. + */ +@SmithyUnstableApi +public record McpPage(List items, Optional> nextPage) { + public McpPage { + items = List.copyOf(items); + } + + public static McpPage last(List items) { + return new McpPage<>(items, Optional.empty()); + } + + public static McpPage continued(List items, NextPage nextPage) { + return new McpPage<>(items, Optional.of(nextPage)); + } + + /** + * Retrieves exactly one subsequent page. + */ + @FunctionalInterface + public interface NextPage { + McpPage fetch(); + } +} 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..b2869e5412 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocol.java @@ -0,0 +1,155 @@ +/* + * 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()); + } + + /** + * Preference used when an initialize request names an unsupported or stateless + * protocol. Higher values are preferred. + * + *

Only initialization-capable, stateful protocols participate. + */ + default int initializationPriority() { + return 0; + } + + 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 Document decorateResult( + Document result, + McpMethod method, + McpServerIdentity serverIdentity, + McpCachePolicy cachePolicy + ) { + return decorateResult(result, method, serverIdentity); + } + + 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..16e88cd0d1 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpProtocolRegistry.java @@ -0,0 +1,205 @@ +/* + * 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(); + var initializationProtocols = this.protocols.values() + .stream() + .filter(protocol -> protocol.supportedMethods().contains(McpMethod.Standard.INITIALIZE)) + .filter(protocol -> !protocol.usesStatelessMetadata()) + .toList(); + var highestPriority = initializationProtocols.stream() + .mapToInt(McpProtocolRegistry::initializationPriority) + .max() + .orElseThrow(() -> new IllegalStateException( + "No initialization-capable MCP protocol is registered")); + var preferred = initializationProtocols.stream() + .filter(protocol -> initializationPriority(protocol) == highestPriority) + .toList(); + if (preferred.size() != 1) { + throw new IllegalStateException( + "Ambiguous initialization fallback priority " + + highestPriority + + " for MCP protocols: " + + preferred.stream() + .map(protocol -> protocol.id().identifier()) + .sorted() + .toList()); + } + initializationFallbackProtocol = preferred.getFirst(); + } + + 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 int initializationPriority(McpProtocol protocol) { + var knownVersion = KnownProtocolVersion.fromIdentifier(protocol.id().identifier()); + return knownVersion == null + ? protocol.initializationPriority() + : knownVersion.ordinal(); + } + + 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..3900d252d0 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRemoteClient.java @@ -0,0 +1,477 @@ +/* + * 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.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Supplier; +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.JsonRpcErrorResponse; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.java.mcp.model.JsonRpcResponse; +import software.amazon.smithy.java.mcp.model.ListPromptsResult; +import software.amazon.smithy.java.mcp.model.ListToolsResult; +import software.amazon.smithy.java.mcp.model.PromptInfo; +import software.amazon.smithy.java.mcp.model.ToolInfo; +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 static final int MAX_LIST_PAGES = 1_000; + + private final AtomicReference> responseNotificationConsumer = new AtomicReference<>(); + private final AtomicReference> requestNotificationConsumer = new AtomicReference<>(); + private final AtomicReference negotiatedProtocol = new AtomicReference<>(); + private final AtomicReference initialization = new AtomicReference<>(); + private final AtomicReference> initializationFlight = new AtomicReference<>(); + private final AtomicLong initializationGeneration = new AtomicLong(); + private final AtomicReference> sessionRecovery = new AtomicReference<>(); + private final ThreadLocal requestProtocol = new ThreadLocal<>(); + + public McpPage listTools() { + return listTools(null, new PageTraversal(), McpMetadata.forProtocol(protocol())); + } + + private McpPage listTools( + String cursor, + PageTraversal traversal, + McpMetadata metadata + ) { + var currentTraversal = traversal.beforeFetch(); + var params = cursor == null + ? null + : Document.of(Map.of("cursor", Document.of(cursor))); + var response = exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.TOOLS_LIST.wireName()) + .id(generateRequestId()) + .jsonrpc("2.0") + .params(metadata.applyTo(params)) + .build()); + requireSuccess(response, "listing tools"); + var result = response.getResult().asShape(ListToolsResult.builder()); + var items = result.getTools().stream().toList(); + onToolsPage(items); + return page( + items, + result.getNextCursor(), + (nextCursor, nextTraversal) -> listTools(nextCursor, nextTraversal, metadata), + currentTraversal); + } + + public McpPage listPrompts() { + return listPrompts(null, new PageTraversal(), McpMetadata.forProtocol(protocol())); + } + + private McpPage listPrompts( + String cursor, + PageTraversal traversal, + McpMetadata metadata + ) { + var currentTraversal = traversal.beforeFetch(); + var params = cursor == null + ? null + : Document.of(Map.of("cursor", Document.of(cursor))); + var response = exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.PROMPTS_LIST.wireName()) + .id(generateRequestId()) + .jsonrpc("2.0") + .params(metadata.applyTo(params)) + .build()); + requireSuccess(response, "listing prompts"); + var result = response.getResult().asShape(ListPromptsResult.builder()); + var items = result.getPrompts().stream().toList(); + return page( + items, + result.getNextCursor(), + (nextCursor, nextTraversal) -> listPrompts(nextCursor, nextTraversal, metadata), + currentTraversal); + } + + protected void onToolsPage(List tools) {} + + final void initialize( + Consumer responseNotificationConsumer, + Consumer requestNotificationConsumer, + JsonRpcRequest initializeRequest, + McpProtocol protocol + ) { + initialize( + responseNotificationConsumer, + requestNotificationConsumer, + initializeRequest, + protocol, + McpProtocolRegistry.create(List.of(), List.of(), false)); + } + + final void initialize( + Consumer responseNotificationConsumer, + Consumer requestNotificationConsumer, + JsonRpcRequest initializeRequest, + McpProtocol protocol, + McpProtocolRegistry protocols + ) { + while (true) { + if (initialized()) { + return; + } + var active = initializationFlight.get(); + if (active != null) { + await(active); + return; + } + + var created = new CompletableFuture(); + if (!initializationFlight.compareAndSet(null, created)) { + continue; + } + try { + if (!initialized()) { + var selected = performInitialize(initializeRequest, protocol, protocols); + this.responseNotificationConsumer.set(responseNotificationConsumer); + this.requestNotificationConsumer.set(requestNotificationConsumer); + initialization.set(new Initialization(initializeRequest, protocol, protocols)); + negotiatedProtocol.set(selected); + exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .jsonrpc("2.0") + .build()); + initializationGeneration.incrementAndGet(); + } + created.complete(null); + return; + } catch (RuntimeException e) { + created.completeExceptionally(e); + throw e; + } finally { + initializationFlight.compareAndSet(created, null); + } + } + } + + private McpProtocol performInitialize( + JsonRpcRequest initializeRequest, + McpProtocol requestedProtocol, + McpProtocolRegistry protocols + ) { + var result = withRequestProtocol( + requestedProtocol, + () -> Objects.requireNonNull( + exchangeForwarded(initializeRequest), + "initialize response")); + requireSuccess(result, "initialization"); + + var negotiatedVersion = McpHttpBinding.stringMember(result.getResult(), "protocolVersion"); + if (negotiatedVersion == null) { + return requestedProtocol; + } + try { + return protocols.require(ProtocolVersion.parse(negotiatedVersion)); + } catch (McpProtocolException e) { + throw new McpRemoteException( + "Remote MCP server negotiated an unsupported protocol version: " + negotiatedVersion, + e); + } + } + + protected final ProtocolVersion protocolVersion() { + return protocol().protocolVersion(); + } + + protected final McpProtocol protocol() { + var requested = requestProtocol.get(); + if (requested != null) { + return requested; + } + var negotiated = negotiatedProtocol.get(); + return negotiated == null + ? BuiltInProtocols.protocol(ProtocolVersion.defaultVersion()) + : negotiated; + } + + final T usingProtocol(McpProtocol requestedProtocol, Supplier operation) { + return withRequestProtocol(requestedProtocol, operation); + } + + private T withRequestProtocol(McpProtocol protocol, Supplier operation) { + var previous = requestProtocol.get(); + requestProtocol.set(protocol); + try { + return operation.get(); + } finally { + if (previous == null) { + requestProtocol.remove(); + } else { + requestProtocol.set(previous); + } + } + } + + protected final long initializationGeneration() { + return initializationGeneration.get(); + } + + protected final boolean restartSession(long observedGeneration) { + var currentInitialization = initialization.get(); + if (currentInitialization == null) { + return false; + } + + while (true) { + if (initializationGeneration.get() != observedGeneration) { + return true; + } + var active = sessionRecovery.get(); + if (active != null) { + await(active); + return true; + } + + var created = new CompletableFuture(); + if (!sessionRecovery.compareAndSet(null, created)) { + continue; + } + try { + if (initializationGeneration.get() == observedGeneration) { + var selected = performInitialize( + currentInitialization.request(), + currentInitialization.requestedProtocol(), + currentInitialization.protocols()); + negotiatedProtocol.set(selected); + exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .jsonrpc("2.0") + .build()); + initializationGeneration.incrementAndGet(); + } + created.complete(null); + return true; + } catch (RuntimeException e) { + created.completeExceptionally(e); + throw e; + } finally { + sessionRecovery.compareAndSet(created, null); + } + } + } + + final boolean initialized() { + return initialization.get() != null; + } + + final McpProtocol negotiatedProtocol() { + return negotiatedProtocol.get(); + } + + final boolean supportsStatelessProtocol(McpProtocol requestedProtocol) { + return usingProtocol(requestedProtocol, () -> { + var response = exchange(JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(generateRequestId()) + .method(McpMethod.Standard.SERVER_DISCOVER.wireName()) + .params(McpMetadata.forProtocol(requestedProtocol).applyTo(null)) + .build()); + if (response == null) { + throw new McpRemoteException("Remote MCP discovery did not return a response"); + } + if (response.getError() != null) { + return false; + } + if (response.getResult() == null) { + throw new McpRemoteException( + "Remote MCP discovery response did not contain a result"); + } + return true; + }); + } + + private void await(CompletableFuture future) { + try { + future.join(); + } catch (CompletionException e) { + if (e.getCause() instanceof RuntimeException runtimeException) { + throw runtimeException; + } + throw e; + } + } + + /** + * Performs one blocking JSON-RPC exchange. Notifications return {@code null}. + */ + protected abstract JsonRpcResponse exchange(JsonRpcRequest request); + + final JsonRpcResponse exchangeForwarded(JsonRpcRequest request) { + var callerId = request.getId(); + if (callerId == null) { + return exchange(request); + } + + var forwarded = JsonRpcRequest.builder() + .jsonrpc(request.getJsonrpc()) + .id(generateRequestId()) + .method(request.getMethod()) + .params(request.getParams()) + .build(); + var response = exchange(forwarded); + if (response == null) { + return null; + } + + var restored = JsonRpcResponse.builder() + .jsonrpc(response.getJsonrpc()) + .id(callerId); + if (response.getError() != null) { + restored.error(response.getError()); + } else if (response.getResult() == null) { + restored.error(JsonRpcErrorResponse.builder() + .code(-32000) + .message("Remote MCP response did not contain a result or error") + .build()); + } else { + restored.result(response.getResult()); + } + return restored.build(); + } + + /** + * 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()); + } + if (response.getResult() == null) { + throw new McpRemoteException("Remote MCP response during " + action + " did not contain a result"); + } + } + + private McpPage page( + List items, + String nextCursor, + PageFetcher fetcher, + PageTraversal traversal + ) { + if (nextCursor == null || nextCursor.isBlank()) { + return McpPage.last(items); + } + var nextTraversal = traversal.record(nextCursor); + return McpPage.continued(items, () -> fetcher.fetch(nextCursor, nextTraversal)); + } + + @FunctionalInterface + private interface PageFetcher { + McpPage fetch(String cursor, PageTraversal traversal); + } + + private record PageTraversal(int pageCount, Set cursors) { + private PageTraversal() { + this(0, Set.of()); + } + + private PageTraversal { + cursors = Set.copyOf(cursors); + } + + PageTraversal beforeFetch() { + if (pageCount >= MAX_LIST_PAGES) { + throw new McpRemoteException( + "Remote MCP listing exceeded the maximum of " + MAX_LIST_PAGES + " pages"); + } + return new PageTraversal(pageCount + 1, cursors); + } + + PageTraversal record(String cursor) { + if (cursors.contains(cursor)) { + throw new McpRemoteException("Remote MCP listing repeated cursor: " + cursor); + } + var updated = new HashSet<>(cursors); + updated.add(cursor); + return new PageTraversal(pageCount, updated); + } + } + + private record Initialization( + JsonRpcRequest request, + McpProtocol requestedProtocol, + McpProtocolRegistry protocols) {} + + 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..cddad15b56 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpRequestDecoder.java @@ -0,0 +1,187 @@ +/* + * 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(); + if (params != null && !isObject(params)) { + throw invalidParams("params must be an object"); + } + 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 24798d051c..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java +++ /dev/null @@ -1,240 +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.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Set; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -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.ListPromptsResult; -import software.amazon.smithy.java.mcp.model.ListToolsResult; -import software.amazon.smithy.java.mcp.model.PromptInfo; -import software.amazon.smithy.java.mcp.model.ToolInfo; -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); - - // Cap list pages so a server that always returns a fresh, advancing cursor fails the call - // instead of looping forever. MCP cursors are opaque and the spec does not guarantee - // termination; at a typical ~30 items/page this bounds a listing at ~30k items. - private static final int MAX_LIST_PAGES = 1000; - - private final AtomicReference> notificationConsumer = new AtomicReference<>(); - private final AtomicReference> requestNotificationConsumer = new AtomicReference<>(); - private final AtomicReference protocolVersion = - new AtomicReference<>(ProtocolVersion.defaultVersion()); - - public List listTools() { - return listPaginated("tools/list", "listing tools", result -> { - ListToolsResult page = result.asShape(ListToolsResult.builder()); - return new Page<>(page.getTools(), page.getNextCursor()); - }); - } - - public List listPrompts() { - return listPaginated("prompts/list", "listing prompts", result -> { - ListPromptsResult page = result.asShape(ListPromptsResult.builder()); - return new Page<>(page.getPrompts(), page.getNextCursor()); - }); - } - - /** - * Maximum number of pages {@link #listTools()} / {@link #listPrompts()} will fetch before - * aborting, a backstop against a server that keeps returning a fresh, advancing cursor and - * never terminates. Subclasses may override to tighten or relax the bound. - */ - protected int maxListPages() { - return MAX_LIST_PAGES; - } - - /** - * Drives MCP cursor pagination for a {@code tools/list}-style method: repeatedly calls - * {@code method}, threading the previous page's {@code nextCursor} back as the {@code cursor} - * request param, and accumulates items across all pages in page order until the server stops - * returning a cursor. A single-page server (no {@code nextCursor}) makes exactly one round-trip. - * - *

Three guards bound a misbehaving server: an absent or blank {@code nextCursor} ends - * pagination; a previously-seen cursor (including a non-advancing {@code A -> B -> A} cycle) - * aborts; and the page count is capped at {@link #maxListPages()}. - */ - private List listPaginated(String method, String errorLabel, PageExtractor extractor) { - List all = new ArrayList<>(); - // Cursors already requested this call, so a repeated or cycling cursor is caught immediately - // rather than only when two identical cursors happen to be adjacent. - Set seenCursors = new HashSet<>(); - String cursor = null; - int page = 0; - do { - if (++page > maxListPages()) { - throw new IllegalStateException( - "Aborting " + method + ": server returned more than " + maxListPages() - + " pages without terminating (possible pagination bug or misbehaving server)"); - } - - JsonRpcRequest.Builder requestBuilder = JsonRpcRequest.builder() - .method(method) - .id(generateRequestId()) - .jsonrpc("2.0"); - if (cursor != null) { - requestBuilder.params(Document.of(Map.of("cursor", Document.of(cursor)))); - } - - JsonRpcResponse response = rpc(requestBuilder.build()).join(); - if (response.getError() != null) { - throw new RuntimeException("Error " + errorLabel + ": " + response.getError().getMessage()); - } - - Document result = response.getResult(); - if (result == null) { - throw new RuntimeException( - "Error " + errorLabel + ": response contained neither a result nor an error"); - } - - Page parsed = extractor.extract(result); - all.addAll(parsed.items()); - - // MCP signals "no more pages" by omitting nextCursor; defensively treat a blank cursor the - // same way, since some servers send "" instead of omitting the field. - String nextCursor = parsed.nextCursor(); - if (nextCursor != null && nextCursor.isBlank()) { - nextCursor = null; - } - if (nextCursor != null && !seenCursors.add(nextCursor)) { - throw new IllegalStateException( - "Aborting " + method + ": server repeated a pagination cursor (no forward progress)"); - } - cursor = nextCursor; - } while (cursor != null); - - LOG.debug("{}: fetched {} item(s) across {} page(s)", method, all.size(), page); - return List.copyOf(all); - } - - /** One page of a paginated list: the page's items plus the server's {@code nextCursor} (null when last). */ - private record Page(List items, String nextCursor) {} - - /** Parses a {@code *_/list} result {@code Document} into its items and {@code nextCursor}. */ - @FunctionalInterface - private interface PageExtractor { - Page extract(Document result); - } - - public void initialize( - 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 a0be512563..0000000000 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java +++ /dev/null @@ -1,1588 +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.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Set; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -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); - - // The tool, prompt, proxy, and service registries are held as immutable snapshots behind volatile - // references (copy-on-write). Readers (tools/list, prompts/list, tools/call dispatch, shutdown) - // read the current snapshot with no locking and always see a complete, consistent map. Every - // mutation (proxy init, dynamic add, and tools/list_changed refresh) rebuilds the affected - // snapshot under registryLock and publishes it atomically, so concurrent mutators cannot lose - // each other's updates and readers never observe a half-updated registry. Network I/O - // (proxy.listTools()/listPrompts()) is always performed outside the lock. - private final Object registryLock = new Object(); - private volatile Map tools; - private volatile Map prompts; - private volatile Map proxies; - private volatile Map services; - private final String serviceName; - private final String version; - private final 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; - // Set once via setNotificationWriter() and read later from the refresh/proxy paths on other - // threads, so it is volatile for safe publication. - private volatile Consumer notificationWriter; - - // Runs tools/list_changed refreshes off the transport's reader thread. A synchronous refresh calls - // listTools() whose response is read by that same reader thread, so doing it inline deadlocks it. - // A single thread also serializes refreshes from different proxies with each other. It is a daemon - // thread that lives for the process; McpService has no explicit lifecycle so it is never shut down. - private final ExecutorService toolRefreshExecutor = - Executors.newSingleThreadExecutor(r -> { - var t = new Thread(r, "mcp-tools-refresh"); - t.setDaemon(true); - return t; - }); - - McpService( - Map services, - List proxyList, - String name, - String version, - ToolFilter toolFilter, - McpMetricsObserver metricsObserver, - McpServerInterceptor interceptor - ) { - // Only services needs copying: it is supplied by the builder, which may still hold or mutate - // it. The tools, prompts, and proxies maps are all built fresh here and never referenced - // again, so snapshot() can wrap them without copying. - this.services = snapshot(new LinkedHashMap<>(services)); - this.schemaIndex = - SchemaIndex.compose(services.values().stream().map(Service::schemaIndex).toArray(SchemaIndex[]::new)); - this.tools = snapshot(createTools(services)); - this.prompts = snapshot(PromptLoader.loadPrompts(services.values())); - this.serviceName = name; - this.version = version; - var proxyMap = new LinkedHashMap(); - for (var proxy : proxyList) { - proxyMap.put(proxy.name(), proxy); - } - this.proxies = snapshot(proxyMap); - this.toolFilter = toolFilter; - this.metricsObserver = metricsObserver; - this.interceptor = interceptor; - } - - /** - * 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 promptValues = prompts.values(); - var promptInfos = new ArrayList(promptValues.size()); - for (var prompt : promptValues) { - promptInfos.add(prompt.promptInfo()); - } - var result = ListPromptsResult.builder() - .prompts(promptInfos) - .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 toolValues = tools.values(); - var toolInfos = new ArrayList(toolValues.size()); - for (var tool : toolValues) { - if (toolFilter.allowTool(tool.serverId(), tool.toolInfo().getName())) { - toolInfos.add(extractToolInfo(tool, protocolVersion)); - } - } - var result = ListToolsResult.builder() - .tools(toolInfos) - .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 -> { - if ("notifications/tools/list_changed".equals(notification.getMethod())) { - LOG.debug("Received tools/list_changed notification from proxy: {}", proxy.name()); - // Refresh on a separate thread. This notification is delivered on the proxy's transport - // reader thread, and refreshProxyTools() calls listTools() whose response is read by that - // same thread, so doing it inline would deadlock the reader. - toolRefreshExecutor.execute(() -> refreshProxyTools(proxy)); - } - // Forward the notification - if (baseNotificationWriter != null) { - baseNotificationWriter.accept(notification); - } - }; - } - - /** - * Re-fetches a proxy's tools after a {@code tools/list_changed} notification and swaps them into the - * registry. Runs off the transport reader thread (see caller). The network fetch happens outside - * {@code registryLock}; only the in-memory snapshot swap is locked. Fetches first so a failed or - * slow refresh never wipes the current tools, then adds the new set before pruning this proxy's - * stale entries, so a concurrent {@code tools/list} never observes a gap (at worst a brief superset). - */ - void refreshProxyTools(McpServerProxy proxy) { - List proxyTools; - try { - proxyTools = proxy.listTools(); - } catch (Exception e) { - LOG.error("Failed to re-fetch tools from proxy: {}", proxy.name(), e); - return; - } - // Fast path: the proxy reports no tools and has none currently registered, so there is - // nothing to add and nothing to prune. Reads the current snapshot lock-free and avoids the - // set allocation, map copy, and lock entirely. (If a tool for this proxy is added - // concurrently right after this check, that add publishes it and a later refresh reconciles.) - if (proxyTools.isEmpty() && !hasToolsFor(proxy)) { - return; - } - synchronized (registryLock) { - Set newNames = new HashSet<>(); - // LinkedHashMap so a refresh keeps the existing order of every other tool: re-put entries - // stay in place and genuinely new tools append, rather than the whole listing reshuffling - // to hash order on each tools/list_changed. - var next = new LinkedHashMap<>(tools); - for (var toolInfo : proxyTools) { - newNames.add(toolInfo.getName()); - next.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy)); - } - next.entrySet() - .removeIf(entry -> entry.getValue().proxy() == proxy && !newNames.contains(entry.getKey())); - tools = snapshot(next); - } - } - - /** Whether any currently registered tool belongs to the given proxy. Reads the current snapshot. */ - private boolean hasToolsFor(McpServerProxy proxy) { - for (var tool : tools.values()) { - if (tool.proxy() == proxy) { - return true; - } - } - return false; - } - - /** - * Publishes a registry snapshot by wrapping the given map unmodifiable. The caller hands off - * ownership: the argument must be a freshly built map that is never mutated or retained after - * this call, since it becomes the live snapshot without being copied. Callers build these as - * {@link LinkedHashMap}s so {@code tools/list} and {@code prompts/list} return a stable, - * deterministic insertion order across refreshes and dynamic additions. - */ - private static Map snapshot(Map ownedMap) { - return Collections.unmodifiableMap(ownedMap); - } - - /** - * Starts proxies without initializing them. - */ - 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); - } - // Isolate each proxy: a failure fetching one proxy's tools or prompts must not abort - // discovery for the rest. - registerProxyListing(proxy); - } - } - } - - /** - * Fetches a proxy's tools and prompts (outside {@code registryLock}) and merges them into the - * registries under the lock. A failure fetching either list is logged and skipped so it cannot - * abort discovery for other proxies. - */ - private void registerProxyListing(McpServerProxy proxy) { - List proxyTools = null; - try { - proxyTools = proxy.listTools(); - } catch (Exception e) { - LOG.error("Failed to fetch tools from proxy: {}", proxy.name(), e); - } - - List proxyPrompts = null; - try { - proxyPrompts = proxy.listPrompts(); - } catch (Exception e) { - LOG.error("Failed to fetch prompts from proxy: {}", proxy.name(), e); - } - - if (proxyTools == null && proxyPrompts == null) { - return; - } - - synchronized (registryLock) { - if (proxyTools != null) { - var nextTools = new LinkedHashMap<>(tools); - for (var toolInfo : proxyTools) { - nextTools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy)); - } - tools = snapshot(nextTools); - } - if (proxyPrompts != null) { - var nextPrompts = new LinkedHashMap<>(prompts); - for (var promptInfo : proxyPrompts) { - var normalizedName = PromptLoader.normalize(promptInfo.getName()); - nextPrompts.putIfAbsent(normalizedName, new Prompt(promptInfo, proxy)); - } - prompts = snapshot(nextPrompts); - } - } - } - - /** - * 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) { - var newTools = createTools(Map.of(id, service)); - synchronized (registryLock) { - var nextServices = new LinkedHashMap<>(services); - nextServices.put(id, service); - services = snapshot(nextServices); - - var nextTools = new LinkedHashMap<>(tools); - nextTools.putAll(newTools); - tools = snapshot(nextTools); - } - } - - public void addNewProxy( - McpServerProxy mcpServerProxy, - Consumer responseWriter - ) { - synchronized (registryLock) { - var nextProxies = new LinkedHashMap<>(proxies); - nextProxies.put(mcpServerProxy.name(), mcpServerProxy); - proxies = snapshot(nextProxies); - } - - mcpServerProxy.start(); - - // Fetches tools/prompts (network I/O) outside the lock, then swaps under it. - registerProxyListing(mcpServerProxy); - } - - /** - * Checks if a service or proxy with the given ID exists. - */ - public boolean containsMcpServer(String id) { - return services.containsKey(id) || proxies.containsKey(id); - } - - /** - * Returns an immutable snapshot of the registered proxies at the time of the call. Subsequent - * additions via {@link #addNewProxy} are not reflected in a previously returned snapshot. - */ - public Map getProxies() { - return proxies; - } - - 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 LinkedHashMap(); - 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..a77a5885a3 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpSources.java @@ -0,0 +1,53 @@ +/* + * 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); + + McpCursorPage listTools(String cursor); + + McpCursorPage listPrompts(String cursor); + + Map remoteClients(); + + boolean containsServer(String id); + + void bindTransport( + Consumer notificationWriter, + Consumer responseWriter + ); + + void initializeRemoteClients(McpProtocol protocol); + + default void ensureRemoteCatalogLoaded() { + ensureRemoteCatalogLoaded(BuiltInProtocols.protocol(ProtocolVersion.defaultVersion())); + } + + void ensureRemoteCatalogLoaded(McpProtocol protocol); + + 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..815a5dee57 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpToolExecutor.java @@ -0,0 +1,168 @@ +/* + * 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; + private final ToolFilter toolFilter; + + McpToolExecutor( + McpSources sources, + McpWireCodec wireCodec, + McpInterceptor interceptor, + McpProtocolRegistry protocols, + ToolFilter toolFilter + ) { + this.sources = sources; + this.wireCodec = wireCodec; + this.interceptor = interceptor; + this.protocols = protocols; + this.toolFilter = toolFilter; + } + + McpOutcome execute(McpCall.CallTool call, McpRequestContext requestContext) { + var descriptor = sources.tool(call.name()); + if (descriptor == null + || !toolFilter.allowTool(descriptor.serverId(), descriptor.info().getName())) { + 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, requestContext); + }; + } + + private McpOutcome invokeLocal( + McpToolDescriptor descriptor, + McpToolDescriptor.LocalTarget target, + McpCall.CallTool call, + McpRequestContext requestContext + ) { + var operation = target.operation(); + var adapter = sources.snapshot().documentAdapter(); + final SerializableShape input; + try { + var inputDocument = adapter.toSmithy( + call.arguments(), + operation.getApiOperation().inputSchema()); + input = inputDocument.asShape(operation.getApiOperation().inputBuilder()); + } catch (RuntimeException e) { + var cause = unwrap(e); + var message = cause.getMessage(); + if (message == null || message.isBlank()) { + message = cause.getClass().getSimpleName(); + } + throw new McpProtocolException( + -32602, + "Invalid arguments for tool " + call.name() + ": " + message); + } + + 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, + McpRequestContext requestContext + ) { + var protocol = protocols.require(requestContext.protocolVersion()); + return target.client() + .usingProtocol( + protocol, + () -> wireCodec.decode(target.client().exchangeForwarded(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..4af58ba461 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpTransportContext.java @@ -0,0 +1,38 @@ +/* + * 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(); + + /** + * Returns whether this transport can deliver unsolicited server messages. + */ + default boolean supportsServerNotifications() { + return false; + } + + record Stdio() implements McpTransportContext { + @Override + public boolean supportsServerNotifications() { + return true; + } + } + + 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..403d9277a2 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpWireCodec.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.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) { + return metadata == null ? params : metadata.applyTo(params); + } +} 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..54b24aace3 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; @@ -75,8 +75,21 @@ public PromptInfo promptInfo() { * @return GetPromptResult with processed template or proxy response */ public GetPromptResult getPromptResult(Document arguments, Document requestId) { + return proxy == null + ? buildLocalPromptResult(arguments) + : getPromptResult(arguments, requestId, McpMetadata.EMPTY, proxy.protocol()); + } + + GetPromptResult getPromptResult( + Document arguments, + Document requestId, + McpMetadata metadata, + McpProtocol protocol + ) { if (proxy != null) { - return delegateToProxy(arguments, requestId); + return proxy.usingProtocol( + protocol, + () -> delegateToProxy(arguments, requestId, metadata)); } return buildLocalPromptResult(arguments); } @@ -84,7 +97,11 @@ public GetPromptResult getPromptResult(Document arguments, Document requestId) { /** * Delegates the prompt request to the proxy server via RPC. */ - private GetPromptResult delegateToProxy(Document arguments, Document requestId) { + private GetPromptResult delegateToProxy( + Document arguments, + Document requestId, + McpMetadata metadata + ) { Map params = new HashMap<>(); params.put("name", Document.of(promptInfo.getName())); if (arguments != null) { @@ -94,16 +111,15 @@ private GetPromptResult delegateToProxy(Document arguments, Document requestId) JsonRpcRequest request = JsonRpcRequest.builder() .method("prompts/get") .id(requestId) - .params(Document.of(params)) + .params(metadata.applyTo(Document.of(params))) .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.exchangeForwarded(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..a35f1b4538 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/SmithyDocumentAdapter.java @@ -0,0 +1,254 @@ +/* + * 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 INTEGER, LONG, BIG_INTEGER, FLOAT, DOUBLE, BIG_DECIMAL -> + Document.of(new BigDecimal(document.asNumber().toString())); + default -> badType(fromType, toType); + }; + case BIG_INTEGER -> switch (fromType) { + case STRING -> Document.of(new BigInteger(document.asString())); + case INTEGER, LONG, BIG_INTEGER -> Document.of(document.asBigInteger()); + 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 63% 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 e6d130a6c1..0b2d6e42df 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 @@ -18,35 +18,34 @@ import java.time.Duration; 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 final Duration requestTimeout; - private StdioProxy(Builder builder) { + private StdioMcpClient(Builder builder) { processBuilder = new ProcessBuilder(); processBuilder.command().add(builder.command); @@ -106,18 +105,18 @@ public Builder workingDirectory(File workingDirectory) { /** * Per-request timeout: a request that never receives a matching response (e.g. a server that * stays alive but goes silent) fails after this duration instead of blocking the caller - * forever. Defaults to 5 minutes, symmetric with {@link HttpMcpProxy}. + * forever. Defaults to 5 minutes, symmetric with {@link HttpMcpClient}. */ public Builder timeout(Duration timeout) { this.timeout = timeout; 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); } } @@ -126,39 +125,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); @@ -167,30 +163,22 @@ 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(); } - // Fail a request that never receives a matching response (server alive but silent) instead of - // blocking the caller forever; symmetric with HttpMcpProxy's request timeout. orTimeout() - // completes responseFuture itself on timeout, so the caller (which holds responseFuture) sees - // the TimeoutException; the derived stage exists only to remove the pending-request entry on - // any completion (success, error, or timeout). Skipped when the write above already failed and - // completed the future. - if (!responseFuture.isDone()) { - responseFuture.orTimeout(requestTimeout.toMillis(), MILLISECONDS) - .whenComplete((response, error) -> pendingRequests.remove(requestId)); + try { + return pending.await(requestTimeout); + } finally { + pendingRequests.remove(requestId, pending); } - - return responseFuture; } - 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()); }; } @@ -233,19 +221,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); } @@ -262,8 +250,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(); } }); @@ -274,51 +262,75 @@ 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(Duration timeout) { + try { + var completed = result.poll(timeout.toMillis(), MILLISECONDS); + if (completed == null) { + throw new McpRemoteException("Timed out waiting for MCP response"); + } + return switch (completed) { + 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..3028a41cf7 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServer.java @@ -0,0 +1,206 @@ +/* + * 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 java.util.concurrent.atomic.AtomicBoolean; +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); + private final AtomicBoolean shuttingDown = new AtomicBoolean(); + + 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) { + if (!shuttingDown.get()) { + 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() { + if (shuttingDown.compareAndSet(false, true)) { + requests.shutdownNow(); + engine.close(); + try { + input.close(); + } catch (IOException e) { + LOG.debug("Error closing MCP input during shutdown", e); + } + listener.interrupt(); + if (listener.getState() == Thread.State.NEW) { + done.countDown(); + } + } + return CompletableFuture.runAsync(() -> { + try { + done.await(); + requests.awaitTermination(30, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("Interrupted while shutting down MCP server", e); + } + }); + } + + 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..419c78f138 --- /dev/null +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioMcpServerBuilder.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.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; + private McpCachePolicy cachePolicy = McpCachePolicy.DEFAULT; + + 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 StdioMcpServerBuilder cachePolicy(McpCachePolicy cachePolicy) { + this.cachePolicy = Objects.requireNonNull(cachePolicy, "cachePolicy"); + 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) + .cachePolicy(cachePolicy) + .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 56% 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 c5c69f298f..311b2a1aa9 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 @@ -12,12 +12,18 @@ import com.sun.net.httpserver.HttpServer; import java.io.IOException; import java.io.OutputStream; +import java.math.BigInteger; import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -38,23 +44,26 @@ 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 ExecutorService mockServerExecutor; + private HttpMcpClient proxy; private String serverUrl; @BeforeEach void setUp() throws IOException { mockServer = HttpServer.create(new InetSocketAddress(0), 0); + mockServerExecutor = Executors.newVirtualThreadPerTaskExecutor(); + mockServer.setExecutor(mockServerExecutor); int port = mockServer.getAddress().getPort(); serverUrl = "http://localhost:" + port + "/mcp"; mockServer.createContext("/mcp", new MockMcpHandler()); mockServer.start(); - proxy = HttpMcpProxy.builder() + proxy = HttpMcpClient.builder() .endpoint(serverUrl) .name("Test MCP") .build(); @@ -65,22 +74,25 @@ void tearDown() { if (mockServer != null) { mockServer.stop(0); } + if (mockServerExecutor != null) { + mockServerExecutor.shutdownNow(); + } 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()) @@ -91,7 +103,7 @@ void testBuilderRejectsSignerAndAuthSchemeTogether() { @Test void testBuilderRejectsAuthSchemeWithoutIdentityResolver() { assertThrows(IllegalArgumentException.class, - () -> HttpMcpProxy.builder() + () -> HttpMcpClient.builder() .endpoint(serverUrl) .authScheme(new TestAuthScheme()) .build()); @@ -100,7 +112,7 @@ void testBuilderRejectsAuthSchemeWithoutIdentityResolver() { @Test void testBuilderRejectsIdentityResolverWithoutAuthScheme() { assertThrows(IllegalArgumentException.class, - () -> HttpMcpProxy.builder() + () -> HttpMcpClient.builder() .endpoint(serverUrl) .identityResolver(TestIdentityResolver.INSTANCE) .build()); @@ -122,7 +134,7 @@ void testAuthSchemeSignsRequest() throws IOException { exchange.close(); }); - HttpMcpProxy authProxy = HttpMcpProxy.builder() + HttpMcpClient authProxy = HttpMcpClient.builder() .endpoint(serverUrl) .authScheme(new TestAuthScheme()) .identityResolver(TestIdentityResolver.INSTANCE) @@ -134,12 +146,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 @@ -161,7 +173,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) @@ -174,28 +186,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(); @@ -207,13 +219,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(); @@ -225,27 +237,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 @@ -256,8 +268,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()); @@ -267,10 +278,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(); + }); + + var notification = JsonRpcRequest.builder() + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .jsonrpc("2.0") + .build(); - ExecutionException exception = assertThrows(ExecutionException.class, future::get); - assertTrue(exception.getCause() instanceof NullPointerException); + assertNull(proxy.exchange(notification)); } @Test @@ -287,12 +312,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")); } @@ -300,7 +325,7 @@ void testRpcHttpError() throws IOException { void testStartAndShutdown() { assertDoesNotThrow(() -> { proxy.start(); - proxy.shutdown().join(); + proxy.close(); }); } @@ -316,8 +341,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()); @@ -325,6 +349,22 @@ void testSseStreamingResponse() throws IOException { assertEquals("final result", response.getResult().asString()); } + @Test + void testSseIgnoresResponsesForOtherRequestIds() throws IOException { + mockServer.removeContext("/mcp"); + mockServer.createContext("/mcp", new SseMismatchedResponseHandler()); + + var response = proxy.exchange(JsonRpcRequest.builder() + .method("test/streaming") + .id(Document.of(BigInteger.ONE)) + .jsonrpc("2.0") + .build()); + + assertNotNull(response); + assertEquals(BigInteger.ONE, response.getId().asBigInteger()); + assertEquals("matching result", response.getResult().asString()); + } + @Test void testSseStreamingWithNotifications() throws IOException { // Track notifications @@ -345,7 +385,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") @@ -353,8 +393,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); @@ -367,6 +406,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 @@ -379,8 +449,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); @@ -401,8 +470,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); @@ -423,8 +491,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); @@ -464,8 +531,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()); @@ -477,13 +543,81 @@ 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()); } + @Test + void expiredSessionIsReinitializedAndTheRequestIsRetried() throws IOException { + var handler = new RecoveringSessionHandler(); + mockServer.removeContext("/mcp"); + mockServer.createContext("/mcp", handler); + + var protocol = BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25); + proxy.initialize( + ignored -> {}, + ignored -> {}, + JsonRpcRequest.builder() + .method(McpMethod.Standard.INITIALIZE.wireName()) + .id(Document.of(1)) + .jsonrpc("2.0") + .params(Document.of(Map.of( + "protocolVersion", + Document.of(protocol.id().identifier())))) + .build(), + protocol); + + var response = proxy.exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.PING.wireName()) + .id(Document.of(2)) + .jsonrpc("2.0") + .build()); + + assertEquals("recovered", response.getResult().asString()); + assertEquals(2, handler.initializeCount.get()); + } + + @Test + void concurrentStaleSessionResponsesShareOneRecovery() throws Exception { + var handler = new ConcurrentRecoveringSessionHandler(); + mockServer.removeContext("/mcp"); + mockServer.createContext("/mcp", handler); + + var protocol = BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25); + proxy.initialize( + ignored -> {}, + ignored -> {}, + JsonRpcRequest.builder() + .method(McpMethod.Standard.INITIALIZE.wireName()) + .id(Document.of(1)) + .jsonrpc("2.0") + .params(Document.of(Map.of( + "protocolVersion", + Document.of(protocol.id().identifier())))) + .build(), + protocol); + + try (var executor = Executors.newVirtualThreadPerTaskExecutor()) { + var first = executor.submit(() -> proxy.exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.PING.wireName()) + .id(Document.of(2)) + .jsonrpc("2.0") + .build())); + var second = executor.submit(() -> proxy.exchange(JsonRpcRequest.builder() + .method(McpMethod.Standard.PING.wireName()) + .id(Document.of(3)) + .jsonrpc("2.0") + .build())); + + assertEquals("recovered", first.get().getResult().asString()); + assertEquals("recovered", second.get().getResult().asString()); + } + assertEquals(2, handler.initializeCount.get()); + assertEquals(2, handler.recoveredRequestCount.get()); + } + private static class SseStreamingHandler implements HttpHandler { @Override public void handle(HttpExchange exchange) throws IOException { @@ -500,9 +634,167 @@ public void handle(HttpExchange exchange) throws IOException { } } + private static final class RecoveringSessionHandler implements HttpHandler { + private final AtomicInteger initializeCount = new AtomicInteger(); + + @Override + public void handle(HttpExchange exchange) throws IOException { + var request = JsonRpcRequest.builder() + .deserialize(JSON_CODEC.createDeserializer(exchange.getRequestBody().readAllBytes())) + .build(); + if (McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName().equals(request.getMethod())) { + exchange.sendResponseHeaders(202, -1); + exchange.close(); + return; + } + + if (McpHttpBinding.isInitialize(request)) { + var count = initializeCount.incrementAndGet(); + exchange.getResponseHeaders().set("Mcp-Session-Id", "session-" + count); + writeResponse( + exchange, + request.getId(), + Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier())))); + return; + } + + var sessionId = exchange.getRequestHeaders().getFirst("Mcp-Session-Id"); + if ("session-1".equals(sessionId)) { + exchange.sendResponseHeaders(404, -1); + exchange.close(); + return; + } + writeResponse(exchange, request.getId(), Document.of("recovered")); + } + + private void writeResponse( + HttpExchange exchange, + Document id, + Document result + ) throws IOException { + var body = JSON_CODEC.serializeToString(JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(id) + .result(result) + .build()).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 final class ConcurrentRecoveringSessionHandler implements HttpHandler { + private final AtomicInteger initializeCount = new AtomicInteger(); + private final AtomicInteger recoveredRequestCount = new AtomicInteger(); + private final AtomicInteger staleRequestOrder = new AtomicInteger(); + private final CountDownLatch staleRequests = new CountDownLatch(2); + private final CountDownLatch firstRequestRecovered = new CountDownLatch(1); + + @Override + public void handle(HttpExchange exchange) throws IOException { + var request = JsonRpcRequest.builder() + .deserialize(JSON_CODEC.createDeserializer(exchange.getRequestBody().readAllBytes())) + .build(); + if (McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName().equals(request.getMethod())) { + exchange.sendResponseHeaders(202, -1); + exchange.close(); + return; + } + + if (McpHttpBinding.isInitialize(request)) { + var count = initializeCount.incrementAndGet(); + exchange.getResponseHeaders().set("Mcp-Session-Id", "session-" + count); + writeResponse( + exchange, + request.getId(), + Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier())))); + return; + } + + var sessionId = exchange.getRequestHeaders().getFirst("Mcp-Session-Id"); + if ("session-1".equals(sessionId)) { + staleRequests.countDown(); + try { + assertTrue(staleRequests.await(5, TimeUnit.SECONDS)); + if (staleRequestOrder.incrementAndGet() > 1) { + assertTrue(firstRequestRecovered.await(5, TimeUnit.SECONDS)); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("interrupted", e); + } + exchange.sendResponseHeaders(404, -1); + exchange.close(); + return; + } + + assertEquals("session-2", sessionId); + recoveredRequestCount.incrementAndGet(); + firstRequestRecovered.countDown(); + writeResponse(exchange, request.getId(), Document.of("recovered")); + } + + private void writeResponse( + HttpExchange exchange, + Document id, + Document result + ) throws IOException { + var body = JSON_CODEC.serializeToString(JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(id) + .result(result) + .build()).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 SseMismatchedResponseHandler implements HttpHandler { + @Override + public void handle(HttpExchange exchange) throws IOException { + var sseResponse = """ + data: {"jsonrpc":"2.0","id":999,"result":"wrong result"} + + data: {"jsonrpc":"2.0","id":1,"result":"matching result"} + + """; + + exchange.getResponseHeaders().set("Content-Type", "text/event-stream"); + var responseBytes = sseResponse.getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, responseBytes.length); + try (var output = exchange.getResponseBody()) { + output.write(responseBytes); + } finally { + exchange.close(); + } + } + } + private static class SseStreamingWithNotificationsHandler implements HttpHandler { @Override public void handle(HttpExchange exchange) throws IOException { + var request = JsonRpcRequest.builder() + .deserialize(JSON_CODEC.createDeserializer(exchange.getRequestBody().readAllBytes())) + .build(); + if (request.getId() == null) { + exchange.sendResponseHeaders(202, -1); + exchange.close(); + return; + } + StringBuilder sseResponse = new StringBuilder(); // Send a notification first @@ -510,7 +802,9 @@ public void handle(HttpExchange exchange) throws IOException { "data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":50}}\n\n"); // Then send the final response - sseResponse.append("data: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":\"final result\"}\n\n"); + sseResponse.append("data: {\"jsonrpc\":\"2.0\",\"id\":") + .append(request.getId().asNumber()) + .append(",\"result\":\"final result\"}\n\n"); exchange.getResponseHeaders().set("Content-Type", "text/event-stream"); byte[] responseBytes = sseResponse.toString().getBytes(StandardCharsets.UTF_8); @@ -524,6 +818,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 { @@ -578,10 +944,9 @@ public void handle(HttpExchange exchange) throws IOException { } @Test - void testListToolsPaginatesAcrossPages() throws IOException { - // Real wire round-trip: the server pages tools/list with a nextCursor, so listTools() must - // follow it and read nextCursor off the actually-deserialized response, not a hand-built - // Document like the McpServerProxy unit tests use. + void testListToolsReturnsAContinuationForTheNextPage() throws IOException { + // Real wire round-trip: the client must expose nextCursor as a lazy continuation and avoid + // fetching the second page until that continuation is invoked. mockServer.removeContext("/mcp"); mockServer.createContext("/mcp", exchange -> { try { @@ -619,9 +984,56 @@ void testListToolsPaginatesAcrossPages() throws IOException { } }); - List tools = proxy.listTools(); + var first = proxy.listTools(); + assertEquals(List.of("t1", "t2"), first.items().stream().map(ToolInfo::getName).toList()); + + var second = first.nextPage().orElseThrow().fetch(); + assertEquals(List.of("t3"), second.items().stream().map(ToolInfo::getName).toList()); + } + + @Test + void modernListingsUseModernHeadersAndMetadata() throws IOException { + var protocolHeader = new AtomicReference(); + var methodHeader = new AtomicReference(); + var requestBody = new AtomicReference(); + mockServer.removeContext("/mcp"); + mockServer.createContext("/mcp", exchange -> { + try { + protocolHeader.set(exchange.getRequestHeaders().getFirst("MCP-Protocol-Version")); + methodHeader.set(exchange.getRequestHeaders().getFirst("Mcp-Method")); + requestBody.set(JsonRpcRequest.builder() + .deserialize(JSON_CODEC.createDeserializer(exchange.getRequestBody().readAllBytes())) + .build()); + + var response = JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(requestBody.get().getId()) + .result(Document.of(ListToolsResult.builder().tools(List.of()).build())) + .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(); + } + }); - assertEquals(List.of("t1", "t2", "t3"), tools.stream().map(ToolInfo::getName).toList()); + proxy.usingProtocol( + BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28), + proxy::listTools); + + assertEquals(KnownProtocolVersion.V2026_07_28.identifier(), protocolHeader.get()); + assertEquals(McpMethod.Standard.TOOLS_LIST.wireName(), methodHeader.get()); + var metadata = requestBody.get().getParams().getMember("_meta"); + assertEquals( + KnownProtocolVersion.V2026_07_28.identifier(), + metadata.getMember(McpWireNames.PROTOCOL_VERSION).asString()); + assertEquals( + Map.of(), + metadata.getMember(McpWireNames.CLIENT_CAPABILITIES).asStringMap()); } private static ToolInfo tool(String name) { 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..1c01ae4bd9 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpArchitectureTest.java @@ -0,0 +1,615 @@ +/* + * 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 extensionProtocolCanBecomeTheInitializationFallback() { + var protocol = protocol( + "2099-01-01", + Set.of(McpMethod.Standard.INITIALIZE), + 100); + try (var engine = McpEngine.builder() + .discoverProtocols(false) + .addProtocol(protocol) + .build()) { + var response = engine.execute( + initializeRequest(new UnknownProtocolVersion("unsupported")), + null); + + assertNull(response.getError()); + assertEquals( + protocol.id().identifier(), + response.getResult().getMember("protocolVersion").asString()); + } + } + + @Test + void ambiguousInitializationFallbackPrioritiesFailConstruction() { + var first = protocol( + "2099-01-01", + Set.of(McpMethod.Standard.INITIALIZE), + 100); + var second = protocol( + "2099-02-01", + Set.of(McpMethod.Standard.INITIALIZE), + 100); + + var error = assertThrows( + IllegalStateException.class, + () -> McpEngine.builder() + .discoverProtocols(false) + .addProtocol(first) + .addProtocol(second) + .build()); + + assertTrue(error.getMessage().contains("Ambiguous initialization fallback priority")); + assertTrue(error.getMessage().contains(first.id().identifier())); + assertTrue(error.getMessage().contains(second.id().identifier())); + } + + @Test + void statelessCacheHintsAreConfigurablePerMethod() { + var cachePolicy = McpCachePolicy.builder() + .hint( + McpMethod.Standard.TOOLS_LIST, + new McpCacheHint(30_000, McpCacheScope.PUBLIC)) + .build(); + try (var engine = McpEngine.builder().cachePolicy(cachePolicy).build()) { + var response = engine.execute( + JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.TOOLS_LIST.wireName()) + .params(statelessParams()) + .build(), + KnownProtocolVersion.V2026_07_28); + + assertNull(response.getError()); + assertEquals(30_000, response.getResult().getMember("ttlMs").asNumber().longValue()); + assertEquals("public", response.getResult().getMember("cacheScope").asString()); + } + } + + @Test + void cacheHintRejectsNegativeTtl() { + assertThrows( + IllegalArgumentException.class, + () -> new McpCacheHint(-1, McpCacheScope.PRIVATE)); + } + + @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 decodeErrorsDoNotProduceResponsesForNotifications() { + try (var engine = McpEngine.builder().build()) { + var notification = JsonRpcRequest.builder() + .jsonrpc("2.0") + .method(McpMethod.Standard.NOTIFICATIONS_INITIALIZED.wireName()) + .params(Document.of("not-an-object")) + .build(); + + assertNull(engine.execute(notification, KnownProtocolVersion.V2025_11_25)); + } + } + + @Test + void metricsObservationToleratesMistypedInitializeMetadata() { + var observer = new McpMetricsObserver() { + @Override + public void onInitialize( + String method, + String extractedProtocolVersion, + boolean rootsListChanged, + boolean sampling, + boolean elicitation, + String clientName, + String clientTitle + ) {} + + @Override + public void onToolCall(String method, String toolName) {} + }; + try (var engine = McpEngine.builder().metricsObserver(observer).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(KnownProtocolVersion.V2025_11_25.identifier()), + "capabilities", + Document.of(Map.of("roots", Document.of("not-an-object"))), + "clientInfo", + Document.of(42)))) + .build(); + + assertNull(engine.execute(request, KnownProtocolVersion.V2025_11_25).getError()); + } + } + + @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 builtInOverrideRetainsTheVersionsInitializationPriority() { + var protocol = protocol( + KnownProtocolVersion.V2025_11_25.identifier(), + Set.of(McpMethod.Standard.INITIALIZE)); + try (var engine = McpEngine.builder() + .discoverProtocols(false) + .overrideProtocol(protocol) + .build()) { + var response = engine.execute( + initializeRequest(new UnknownProtocolVersion("unsupported")), + null); + + assertNull(response.getError()); + assertEquals( + KnownProtocolVersion.V2025_11_25.identifier(), + response.getResult().getMember("protocolVersion").asString()); + } + } + + @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 protocol(identifier, methods, 0); + } + + private TestProtocol protocol( + String identifier, + Set methods, + int initializationPriority + ) { + return new TestProtocol( + McpProtocolId.of(identifier), + methods, + initializationPriority); + } + + 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, + int initializationPriority) 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), + 0)); + } + } + + 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..a7d4c84644 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpCatalogTest.java @@ -0,0 +1,780 @@ +/* + * 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.assertNotNull; +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.CopyOnWriteArrayList; +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.AtomicLong; +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.JsonRpcErrorResponse; +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 McpPage listTools() { + throw new McpRemoteException("unavailable"); + } + }; + var healthy = new TestRemoteClient("healthy") { + @Override + public McpPage listTools() { + return McpPage.last(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 toolPagesPassThroughWithoutAggregatingFutureListings() { + var continuationCalls = new AtomicInteger(); + var remote = new TestRemoteClient("paged") { + @Override + public McpPage listTools() { + return McpPage.continued( + List.of(tool("FirstTool")), + () -> { + continuationCalls.incrementAndGet(); + return McpPage.last(List.of(tool("SecondTool"))); + }); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.ensureRemoteCatalogLoaded(); + + var first = catalog.listTools(null); + assertEquals(List.of("FirstTool"), toolNames(first)); + assertNotNull(first.nextCursor()); + assertEquals(0, continuationCalls.get()); + + var second = catalog.listTools(first.nextCursor()); + assertEquals(List.of("SecondTool"), toolNames(second)); + assertNull(second.nextCursor()); + assertEquals(1, continuationCalls.get()); + assertNotNull(catalog.tool("SecondTool")); + + var fresh = catalog.listTools(null); + assertEquals(List.of("FirstTool"), toolNames(fresh)); + assertNotNull(fresh.nextCursor()); + assertEquals(1, continuationCalls.get()); + + var error = assertThrows( + McpProtocolException.class, + () -> catalog.listTools(first.nextCursor())); + assertEquals(-32602, error.code()); + } + } + + @Test + void promptPagesPassThroughWithoutAggregatingFutureListings() { + var continuationCalls = new AtomicInteger(); + var remote = new TestRemoteClient("paged") { + @Override + public McpPage listPrompts() { + return McpPage.continued( + List.of(prompt("FirstPrompt")), + () -> { + continuationCalls.incrementAndGet(); + return McpPage.last(List.of(prompt("SecondPrompt"))); + }); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.ensureRemoteCatalogLoaded(); + + var first = catalog.listPrompts(null); + assertEquals(List.of("FirstPrompt"), promptNames(first)); + assertNotNull(first.nextCursor()); + assertEquals(0, continuationCalls.get()); + + var second = catalog.listPrompts(first.nextCursor()); + assertEquals(List.of("SecondPrompt"), promptNames(second)); + assertNull(second.nextCursor()); + assertEquals(1, continuationCalls.get()); + assertNotNull(catalog.prompt(PromptLoader.normalize("SecondPrompt"))); + + var fresh = catalog.listPrompts(null); + assertEquals(List.of("FirstPrompt"), promptNames(fresh)); + assertNotNull(fresh.nextCursor()); + assertEquals(1, continuationCalls.get()); + } + } + + @Test + void failedPageFetchDoesNotConsumeTheCursor() { + var continuationCalls = new AtomicInteger(); + var remote = new TestRemoteClient("retryable-page") { + @Override + public McpPage listTools() { + return McpPage.continued( + List.of(tool("FirstTool")), + () -> { + if (continuationCalls.incrementAndGet() == 1) { + throw new McpRemoteException("temporary failure"); + } + return McpPage.last(List.of(tool("SecondTool"))); + }); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.ensureRemoteCatalogLoaded(); + var first = catalog.listTools(null); + + assertThrows( + McpRemoteException.class, + () -> catalog.listTools(first.nextCursor())); + assertEquals( + List.of("SecondTool"), + toolNames(catalog.listTools(first.nextCursor()))); + } + } + + @Test + void refreshRetainsAdvertisedPagesAndExistingCursors() throws Exception { + var refreshComplete = new CountDownLatch(1); + var listings = new AtomicInteger(); + var remote = new TestRemoteClient("refreshing-pages") { + @Override + public McpPage listTools() { + if (listings.incrementAndGet() == 1) { + return McpPage.continued( + List.of(tool("FirstTool")), + () -> McpPage.last(List.of(tool("SecondTool")))); + } + refreshComplete.countDown(); + return McpPage.last(List.of(tool("RefreshedFirstTool"))); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.bindTransport(ignored -> {}, ignored -> {}); + catalog.initializeRemoteClients( + BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25)); + + var first = catalog.listTools(null); + catalog.listTools(first.nextCursor()); + var cursorCreatedBeforeRefresh = catalog.listTools(null).nextCursor(); + + remote.sendNotification(JsonRpcRequest.builder() + .jsonrpc("2.0") + .method(McpMethod.Standard.NOTIFICATIONS_TOOLS_LIST_CHANGED.wireName()) + .build()); + assertTrue(refreshComplete.await(2, SECONDS)); + + assertNotNull(catalog.tool("SecondTool")); + assertEquals( + List.of("SecondTool"), + toolNames(catalog.listTools(cursorCreatedBeforeRefresh))); + } + } + + @Test + void cursorRegistryIsBounded() { + var remote = new TestRemoteClient("bounded-cursors") { + @Override + public McpPage listTools() { + return McpPage.continued( + List.of(tool("FirstTool")), + () -> McpPage.last(List.of(tool("SecondTool")))); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.ensureRemoteCatalogLoaded(); + var oldest = catalog.listTools(null).nextCursor(); + for (int index = 0; index < 1_024; index++) { + catalog.listTools(null); + } + + var error = assertThrows( + McpProtocolException.class, + () -> catalog.listTools(oldest)); + assertEquals(-32602, error.code()); + } + } + + @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 remoteCatalogCanReloadUsingANewerProtocol() { + var versions = new CopyOnWriteArrayList(); + var remote = new TestRemoteClient("versioned") { + @Override + public McpPage listTools() { + versions.add(protocolVersion()); + return McpPage.last(List.of()); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.ensureRemoteCatalogLoaded( + BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25)); + catalog.ensureRemoteCatalogLoaded( + BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28)); + + assertEquals( + List.of( + KnownProtocolVersion.V2025_11_25, + KnownProtocolVersion.V2026_07_28), + versions); + } + } + + @Test + void statelessFrontendInitializesALegacyRemoteWithProxyIdentity() { + var initializeCount = new AtomicInteger(); + var clientName = new AtomicReference(); + var initialized = new AtomicBoolean(); + var remote = new TestRemoteClient("legacy-only") { + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + if (McpMethod.Standard.SERVER_DISCOVER.wireName().equals(request.getMethod())) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .error(JsonRpcErrorResponse.builder() + .code(-32601) + .message("Method not found") + .build()) + .build(); + } + if (McpHttpBinding.isInitialize(request)) { + initializeCount.incrementAndGet(); + clientName.set(request.getParams() + .getMember("clientInfo") + .getMember("name") + .asString()); + initialized.set(true); + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .result(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier())))) + .build(); + } + return super.exchange(request); + } + + @Override + public McpPage listTools() { + if (!initialized.get()) { + throw new McpRemoteException("not initialized"); + } + return McpPage.last(List.of(tool("LegacyTool"))); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.ensureRemoteCatalogLoaded( + BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28)); + + assertEquals(1, initializeCount.get()); + assertEquals("mcp-server", clientName.get()); + assertNotNull(catalog.tool("LegacyTool")); + } + } + + @Test + void failedRemoteInitializationRetriesAfterTheCooldown() { + var initializeCount = new AtomicInteger(); + var initialized = new AtomicBoolean(); + var nanoTime = new AtomicLong(); + var remote = new TestRemoteClient("retryable-initialize") { + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + if (McpMethod.Standard.SERVER_DISCOVER.wireName().equals(request.getMethod())) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .error(JsonRpcErrorResponse.builder() + .code(-32601) + .message("Method not found") + .build()) + .build(); + } + if (McpHttpBinding.isInitialize(request)) { + if (initializeCount.incrementAndGet() == 1) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .error(JsonRpcErrorResponse.builder() + .code(-32000) + .message("temporarily unavailable") + .build()) + .build(); + } + initialized.set(true); + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .result(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_11_25.identifier())))) + .build(); + } + return super.exchange(request); + } + + @Override + public McpPage listTools() { + if (!initialized.get()) { + throw new McpRemoteException("not initialized"); + } + return McpPage.last(List.of(tool("RecoveredTool"))); + } + }; + var modern = BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28); + + try (var catalog = catalogWithClock(remote, nanoTime)) { + catalog.ensureRemoteCatalogLoaded(modern); + assertNull(catalog.tool("RecoveredTool")); + + catalog.ensureRemoteCatalogLoaded(modern); + assertEquals(1, initializeCount.get()); + + nanoTime.set(SECONDS.toNanos(31)); + catalog.ensureRemoteCatalogLoaded(modern); + + assertEquals(2, initializeCount.get()); + assertNotNull(catalog.tool("RecoveredTool")); + } + } + + @Test + void concurrentProtocolLoadsInitializeARemoteOnlyOnce() throws Exception { + var initializeCount = new AtomicInteger(); + var secondInitialize = new CountDownLatch(1); + var remote = new TestRemoteClient("single-flight-initialize") { + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + if (McpHttpBinding.isInitialize(request)) { + if (initializeCount.incrementAndGet() > 1) { + secondInitialize.countDown(); + } + try { + secondInitialize.await(250, java.util.concurrent.TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("interrupted", e); + } + var requestedVersion = request.getParams() + .getMember("protocolVersion") + .asString(); + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .result(Document.of(Map.of( + "protocolVersion", + Document.of(requestedVersion)))) + .build(); + } + return super.exchange(request); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote)); + var executor = Executors.newVirtualThreadPerTaskExecutor()) { + var start = new CountDownLatch(1); + var first = executor.submit(() -> { + start.await(); + catalog.ensureRemoteCatalogLoaded( + BuiltInProtocols.protocol(KnownProtocolVersion.V2024_11_05)); + return null; + }); + var second = executor.submit(() -> { + start.await(); + catalog.ensureRemoteCatalogLoaded( + BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25)); + return null; + }); + + start.countDown(); + first.get(); + second.get(); + + assertEquals(1, initializeCount.get()); + } + } + + @Test + void failedRemoteLoadUsesCooldownInsteadOfBlockingEveryRequest() { + var attempts = new AtomicInteger(); + var nanoTime = new AtomicLong(); + var remote = new TestRemoteClient("down") { + @Override + public McpPage listTools() { + attempts.incrementAndGet(); + throw new McpRemoteException("unavailable"); + } + + @Override + public McpPage listPrompts() { + attempts.incrementAndGet(); + throw new McpRemoteException("unavailable"); + } + }; + + try (var catalog = catalogWithClock(remote, nanoTime)) { + catalog.ensureRemoteCatalogLoaded(); + assertEquals(2, attempts.get()); + + assertTimeoutPreemptively( + Duration.ofMillis(250), + () -> { + catalog.ensureRemoteCatalogLoaded(); + }); + assertEquals(2, attempts.get()); + + nanoTime.set(SECONDS.toNanos(31)); + catalog.ensureRemoteCatalogLoaded(); + assertEquals(4, attempts.get()); + } + } + + @Test + void remoteIoDoesNotBlockCatalogReads() throws Exception { + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var remote = new TestRemoteClient("blocking") { + @Override + public McpPage listTools() { + entered.countDown(); + try { + assertTrue(release.await(5, SECONDS)); + return McpPage.last(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 McpPage listTools() { + observedVersion.set(protocolVersion()); + return McpPage.last(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 McpPage 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 McpPage.last(List.of(tool("notifying-tool"))); + } + }; + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + try { + catalog.bindTransport(ignored -> {}, ignored -> {}); + catalog.initializeRemoteClients( + 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(); + } + } + } + + @Test + void remoteNotificationsFanOutToEveryBoundTransport() { + var first = new AtomicInteger(); + var second = new AtomicInteger(); + var remote = new TestRemoteClient("notifications"); + + try (var catalog = new McpCatalog(Map.of(), List.of(remote))) { + catalog.bindTransport(ignored -> first.incrementAndGet(), ignored -> {}); + catalog.bindTransport(ignored -> second.incrementAndGet(), ignored -> {}); + catalog.initializeRemoteClients( + BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25)); + + remote.sendNotification(JsonRpcRequest.builder() + .jsonrpc("2.0") + .method("notifications/progress") + .build()); + + assertEquals(1, first.get()); + assertEquals(1, second.get()); + } + } + + private TestRemoteClient blockingToolClient( + String name, + String toolName, + CountDownLatch entered, + AtomicBoolean timedOut + ) { + return new TestRemoteClient(name) { + @Override + public McpPage listTools() { + entered.countDown(); + try { + if (!entered.await(2, SECONDS)) { + timedOut.set(true); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new McpRemoteException("interrupted", e); + } + return McpPage.last(List.of(tool(toolName))); + } + }; + } + + private ToolInfo tool(String name) { + return ToolInfo.builder() + .name(name) + .inputSchema(JsonObjectSchema.builder().build()) + .build(); + } + + private PromptInfo prompt(String name) { + return PromptInfo.builder().name(name).build(); + } + + private List toolNames(McpCursorPage page) { + return page.items().stream().map(tool -> tool.info().getName()).toList(); + } + + private List promptNames(McpCursorPage page) { + return page.items() + .stream() + .map(prompt -> prompt.prompt().promptInfo().getName()) + .toList(); + } + + 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 McpCatalog catalogWithClock( + McpRemoteClient remote, + AtomicLong nanoTime + ) { + return new McpCatalog( + Map.of(), + List.of(remote), + McpProtocolRegistry.create(List.of(), List.of(), false), + new McpServerIdentity("mcp-server", "1.0.0"), + nanoTime::get, + SECONDS.toNanos(30)); + } + + 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 McpPage listTools() { + return McpPage.last(List.of()); + } + + @Override + public McpPage listPrompts() { + return McpPage.last(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..97c7dfef86 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpHttpHandlerTest.java @@ -0,0 +1,265 @@ +/* + * 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.assertFalse; +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 java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcErrorResponse; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.java.mcp.model.JsonRpcResponse; + +class McpHttpHandlerTest { + + @Test + void constructingAHandlerDoesNotStartRemoteClients() { + var starts = new AtomicInteger(); + var remote = new McpRemoteClient() { + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + return null; + } + + @Override + public void start() { + starts.incrementAndGet(); + } + + @Override + public void close() {} + + @Override + public String name() { + return "lazy"; + } + }; + + try (var engine = McpEngine.builder().addRemoteClient(remote).build()) { + new McpHttpHandler(engine); + assertEquals(0, starts.get()); + } + } + + @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 httpInitializeDoesNotAdvertiseUnavailableCapabilities() { + try (var engine = McpEngine.builder().build()) { + var response = new McpHttpHandler(engine).handle(initializeRequest(), Map.of()); + var capabilities = response.body().getResult().getMember("capabilities"); + + assertNull(capabilities.getMember("completions")); + assertNull(capabilities.getMember("logging")); + assertFalse(capabilities.getMember("tools").getMember("listChanged").asBoolean()); + assertFalse(capabilities.getMember("prompts").getMember("listChanged").asBoolean()); + } + } + + @Test + void parseAndInvalidRequestErrorsUseBadRequestStatus() { + for (var code : List.of(-32700, -32600)) { + var response = JsonRpcResponse.builder() + .jsonrpc("2.0") + .error(JsonRpcErrorResponse.builder() + .code(code) + .message("bad request") + .build()) + .build(); + + assertEquals(400, McpHttpBinding.statusCode(response, false, false)); + assertEquals(400, McpHttpBinding.statusCode(response, true, false)); + } + } + + @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 nonObjectParamsReturnInvalidParamsInsteadOfEscapingTheHandler() { + try (var engine = McpEngine.builder().build()) { + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.PING.wireName()) + .params(Document.of("not-an-object")) + .build(); + + var response = new McpHttpHandler(engine).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/McpRemoteClientTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpRemoteClientTest.java new file mode 100644 index 0000000000..9e54eac724 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpRemoteClientTest.java @@ -0,0 +1,341 @@ +/* + * 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +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.JsonRpcErrorResponse; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; +import software.amazon.smithy.java.mcp.model.JsonRpcResponse; +import software.amazon.smithy.java.mcp.model.ListPromptsResult; +import software.amazon.smithy.java.mcp.model.ListToolsResult; +import software.amazon.smithy.java.mcp.model.PromptInfo; +import software.amazon.smithy.java.mcp.model.ToolInfo; + +class McpRemoteClientTest { + + private static final class FakeClient extends McpRemoteClient { + private final List requests = new ArrayList<>(); + private final List responses; + private int index; + + FakeClient(List responses) { + this.responses = responses; + } + + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + requests.add(request); + return responses.get(index++); + } + + @Override + public void start() {} + + @Override + public void close() {} + + @Override + public String name() { + return "fake"; + } + } + + @Test + void toolsPaginationFetchesOnePageAtATime() { + var client = new FakeClient(List.of( + toolsResponse(List.of(tool("a"), tool("b")), "CURSOR1"), + toolsResponse(List.of(tool("c"), tool("d")), "CURSOR2"), + toolsResponse(List.of(tool("e")), null))); + + var first = client.listTools(); + assertEquals(List.of("a", "b"), names(first.items())); + assertEquals(1, client.requests.size()); + assertNull(client.requests.getFirst().getParams()); + + var second = first.nextPage().orElseThrow().fetch(); + assertEquals(List.of("c", "d"), names(second.items())); + assertEquals(2, client.requests.size()); + assertEquals("CURSOR1", client.requests.get(1).getParams().getMember("cursor").asString()); + + var third = second.nextPage().orElseThrow().fetch(); + assertEquals(List.of("e"), names(third.items())); + assertFalse(third.nextPage().isPresent()); + assertEquals("CURSOR2", client.requests.get(2).getParams().getMember("cursor").asString()); + } + + @Test + void promptsPaginationFetchesOnePageAtATime() { + var client = new FakeClient(List.of( + promptsResponse(List.of(prompt("p1")), "PC1"), + promptsResponse(List.of(prompt("p2"), prompt("p3")), null))); + + var first = client.listPrompts(); + assertEquals(List.of("p1"), first.items().stream().map(PromptInfo::getName).toList()); + assertEquals(1, client.requests.size()); + + var second = first.nextPage().orElseThrow().fetch(); + assertEquals(List.of("p2", "p3"), second.items().stream().map(PromptInfo::getName).toList()); + assertEquals("PC1", client.requests.get(1).getParams().getMember("cursor").asString()); + } + + @Test + void forwardedRequestsUseClientOwnedIdsAndRestoreTheCallerId() { + var client = new FakeClient(List.of( + success(Document.of(Map.of())), + success(Document.of(Map.of())))); + var callerId = Document.of("shared-caller-id"); + var request = JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(callerId) + .method(McpMethod.Standard.PING.wireName()) + .build(); + + var first = client.exchangeForwarded(request); + var second = client.exchangeForwarded(request); + + assertEquals(callerId, first.getId()); + assertEquals(callerId, second.getId()); + assertFalse(client.requests.getFirst().getId().equals(callerId)); + assertFalse(client.requests.get(1).getId().equals(callerId)); + assertFalse(client.requests.getFirst().getId().equals(client.requests.get(1).getId())); + } + + @Test + void initializeUsesTheProtocolVersionNegotiatedByTheRemote() { + var client = new McpRemoteClient() { + @Override + protected JsonRpcResponse exchange(JsonRpcRequest request) { + if (McpHttpBinding.isInitialize(request)) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(request.getId()) + .result(Document.of(Map.of( + "protocolVersion", + Document.of(KnownProtocolVersion.V2025_03_26.identifier())))) + .build(); + } + return null; + } + + ProtocolVersion currentProtocolVersion() { + return protocolVersion(); + } + + @Override + public void start() {} + + @Override + public void close() {} + + @Override + public String name() { + return "negotiating"; + } + }; + var requested = BuiltInProtocols.protocol(KnownProtocolVersion.V2025_11_25); + + client.initialize( + ignored -> {}, + ignored -> {}, + JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method(McpMethod.Standard.INITIALIZE.wireName()) + .params(Document.of(Map.of( + "protocolVersion", + Document.of(requested.id().identifier())))) + .build(), + requested); + + assertEquals(KnownProtocolVersion.V2025_03_26, client.currentProtocolVersion()); + assertEquals( + KnownProtocolVersion.V2025_11_25, + client.usingProtocol(requested, client::currentProtocolVersion)); + } + + @Test + void modernProtocolMetadataIsRetainedAcrossToolPages() { + var client = new FakeClient(List.of( + toolsResponse(List.of(tool("FirstTool")), "next"), + toolsResponse(List.of(tool("SecondTool")), null))); + var modern = BuiltInProtocols.protocol(KnownProtocolVersion.V2026_07_28); + + var first = client.usingProtocol(modern, client::listTools); + var second = first.nextPage().orElseThrow().fetch(); + + assertEquals(List.of("FirstTool"), names(first.items())); + assertEquals(List.of("SecondTool"), names(second.items())); + for (var request : client.requests) { + var metadata = request.getParams().getMember("_meta"); + assertEquals( + KnownProtocolVersion.V2026_07_28.identifier(), + metadata.getMember(McpWireNames.PROTOCOL_VERSION).asString()); + assertEquals( + Map.of(), + metadata.getMember(McpWireNames.CLIENT_CAPABILITIES).asStringMap()); + } + } + + @Test + void blankCursorEndsListing() { + var client = new FakeClient(List.of(toolsResponse(List.of(tool("a")), ""))); + + var page = client.listTools(); + + assertFalse(page.nextPage().isPresent()); + assertEquals(1, client.requests.size()); + } + + @Test + void laterPageErrorsAreDeferredUntilContinuationIsInvoked() { + var client = new FakeClient(List.of( + toolsResponse(List.of(tool("a")), "next"), + errorResponse("boom"))); + + var first = client.listTools(); + assertEquals(List.of("a"), names(first.items())); + + var error = assertThrows( + McpRemoteException.class, + () -> first.nextPage().orElseThrow().fetch()); + assertTrue(error.getMessage().contains("boom")); + } + + @Test + void repeatedCursorFailsWhenTheContinuationIsInvoked() { + var client = new FakeClient(List.of( + toolsResponse(List.of(tool("FirstTool")), "same"), + toolsResponse(List.of(tool("SecondTool")), "same"))); + + var first = client.listTools(); + + var error = assertThrows( + McpRemoteException.class, + () -> first.nextPage().orElseThrow().fetch()); + assertTrue(error.getMessage().contains("repeated cursor")); + assertEquals(2, client.requests.size()); + } + + @Test + void cachedContinuationCanStartIndependentPaginationTraversals() { + var client = new FakeClient(List.of( + toolsResponse(List.of(tool("FirstTool")), "cursor1"), + toolsResponse(List.of(tool("SecondTool")), "cursor2"), + toolsResponse(List.of(tool("SecondTool")), "cursor2"))); + + var first = client.listTools(); + var continuation = first.nextPage().orElseThrow(); + + var firstBranch = continuation.fetch(); + var secondBranch = continuation.fetch(); + + assertEquals(List.of("SecondTool"), names(firstBranch.items())); + assertEquals(List.of("SecondTool"), names(secondBranch.items())); + assertTrue(firstBranch.nextPage().isPresent()); + assertTrue(secondBranch.nextPage().isPresent()); + assertEquals("cursor1", client.requests.get(1).getParams().getMember("cursor").asString()); + assertEquals("cursor1", client.requests.get(2).getParams().getMember("cursor").asString()); + } + + @Test + void listingStopsBeforeFetchingMoreThanThePageLimit() { + var responses = IntStream.range(0, 1_000) + .mapToObj(index -> toolsResponse( + List.of(tool("Tool" + index)), + "cursor" + index)) + .toList(); + var client = new FakeClient(responses); + + var page = client.listTools(); + for (int index = 1; index < 1_000; index++) { + page = page.nextPage().orElseThrow().fetch(); + } + + var lastPage = page; + var error = assertThrows( + McpRemoteException.class, + () -> lastPage.nextPage().orElseThrow().fetch()); + assertTrue(error.getMessage().contains("maximum of 1000 pages")); + assertEquals(1_000, client.requests.size()); + } + + @Test + void missingResultProducesAnActionableRemoteError() { + var client = new FakeClient(List.of(JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .build())); + + var error = assertThrows(McpRemoteException.class, client::listTools); + + assertTrue(error.getMessage().contains("listing tools")); + assertTrue(error.getMessage().contains("did not contain a result")); + } + + @Test + void pageItemsAreImmutable() { + var client = new FakeClient(List.of(toolsResponse(List.of(tool("a")), null))); + + var page = client.listTools(); + + assertThrows(UnsupportedOperationException.class, () -> page.items().add(tool("b"))); + } + + private static List names(List tools) { + return tools.stream().map(ToolInfo::getName).toList(); + } + + private static ToolInfo tool(String name) { + return ToolInfo.builder().name(name).build(); + } + + private static PromptInfo prompt(String name) { + return PromptInfo.builder().name(name).build(); + } + + private static JsonRpcResponse toolsResponse(List tools, String nextCursor) { + var result = ListToolsResult.builder().tools(tools); + if (nextCursor != null) { + result.nextCursor(nextCursor); + } + return success(Document.of(result.build())); + } + + private static JsonRpcResponse promptsResponse(List prompts, String nextCursor) { + var result = ListPromptsResult.builder().prompts(prompts); + if (nextCursor != null) { + result.nextCursor(nextCursor); + } + return success(Document.of(result.build())); + } + + private static JsonRpcResponse success(Document result) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .result(result) + .build(); + } + + private static JsonRpcResponse errorResponse(String message) { + return JsonRpcResponse.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .error(JsonRpcErrorResponse.builder().code(-32000).message(message).build()) + .build(); + } +} 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/McpServerProxyTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerProxyTest.java deleted file mode 100644 index cc31127886..0000000000 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerProxyTest.java +++ /dev/null @@ -1,245 +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 org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import org.junit.jupiter.api.Test; -import software.amazon.smithy.java.core.serde.document.Document; -import software.amazon.smithy.java.mcp.model.JsonRpcErrorResponse; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; -import software.amazon.smithy.java.mcp.model.JsonRpcResponse; -import software.amazon.smithy.java.mcp.model.ListPromptsResult; -import software.amazon.smithy.java.mcp.model.ListToolsResult; -import software.amazon.smithy.java.mcp.model.PromptInfo; -import software.amazon.smithy.java.mcp.model.ToolInfo; - -class McpServerProxyTest { - - /** - * Test proxy that replays a fixed list of canned responses and records every request it received, - * so pagination behaviour (nextCursor to cursor round-tripping) can be asserted. - */ - private static final class FakeProxy extends McpServerProxy { - private final List requests = new ArrayList<>(); - private final List responses; - private int index = 0; - - FakeProxy(List responses) { - this.responses = responses; - } - - @Override - protected CompletableFuture rpc(JsonRpcRequest request) { - requests.add(request); - return CompletableFuture.completedFuture(responses.get(index++)); - } - - @Override - protected void start() {} - - @Override - protected CompletableFuture shutdown() { - return CompletableFuture.completedFuture(null); - } - - @Override - public String name() { - return "fake"; - } - } - - private static ToolInfo tool(String name) { - return ToolInfo.builder().name(name).build(); - } - - private static PromptInfo prompt(String name) { - return PromptInfo.builder().name(name).build(); - } - - private static JsonRpcResponse toolsResponse(List tools, String nextCursor) { - var result = ListToolsResult.builder().tools(tools); - if (nextCursor != null) { - result.nextCursor(nextCursor); - } - return JsonRpcResponse.builder() - .jsonrpc("2.0") - .id(Document.of(1)) - .result(Document.of(result.build())) - .build(); - } - - private static JsonRpcResponse promptsResponse(List prompts, String nextCursor) { - var result = ListPromptsResult.builder().prompts(prompts); - if (nextCursor != null) { - result.nextCursor(nextCursor); - } - return JsonRpcResponse.builder() - .jsonrpc("2.0") - .id(Document.of(1)) - .result(Document.of(result.build())) - .build(); - } - - private static JsonRpcResponse errorResponse(String message) { - return JsonRpcResponse.builder() - .jsonrpc("2.0") - .id(Document.of(1)) - .error(JsonRpcErrorResponse.builder().code(-32000).message(message).build()) - .build(); - } - - private static JsonRpcResponse emptyResponse() { - // A malformed response carrying neither a result nor an error. - return JsonRpcResponse.builder() - .jsonrpc("2.0") - .id(Document.of(1)) - .build(); - } - - @Test - void listToolsFollowsNextCursorAcrossPages() { - var proxy = new FakeProxy(List.of( - toolsResponse(List.of(tool("a"), tool("b")), "CURSOR1"), - toolsResponse(List.of(tool("c"), tool("d")), "CURSOR2"), - toolsResponse(List.of(tool("e")), null))); - - var tools = proxy.listTools(); - - assertEquals(List.of("a", "b", "c", "d", "e"), - tools.stream().map(ToolInfo::getName).toList()); - assertEquals(3, proxy.requests.size()); - // First page carries no cursor. - assertNull(proxy.requests.get(0).getParams()); - // Each subsequent page echoes the prior page's nextCursor as the cursor param. - assertEquals("CURSOR1", proxy.requests.get(1).getParams().getMember("cursor").asString()); - assertEquals("CURSOR2", proxy.requests.get(2).getParams().getMember("cursor").asString()); - } - - @Test - void listToolsSinglePageMakesOneCall() { - var proxy = new FakeProxy(List.of( - toolsResponse(List.of(tool("only")), null))); - - var tools = proxy.listTools(); - - assertEquals(1, tools.size()); - assertEquals(1, proxy.requests.size()); - assertNull(proxy.requests.get(0).getParams()); - } - - @Test - void listPromptsFollowsNextCursorAcrossPages() { - var proxy = new FakeProxy(List.of( - promptsResponse(List.of(prompt("p1")), "PC1"), - promptsResponse(List.of(prompt("p2"), prompt("p3")), null))); - - var prompts = proxy.listPrompts(); - - assertEquals(List.of("p1", "p2", "p3"), - prompts.stream().map(PromptInfo::getName).toList()); - assertEquals(2, proxy.requests.size()); - assertEquals("PC1", proxy.requests.get(1).getParams().getMember("cursor").asString()); - } - - @Test - void listToolsAbortsOnRepeatedCursor() { - var proxy = new FakeProxy(List.of( - toolsResponse(List.of(tool("a")), "SAME"), - toolsResponse(List.of(tool("b")), "SAME"))); - - assertThrows(IllegalStateException.class, proxy::listTools); - } - - @Test - void listToolsAbortsAtPageCap() { - // A server that always advances the cursor never trips the repeated-cursor guard, so the - // MAX_LIST_PAGES cap must stop it. Supply 1001 ever-advancing pages; only 1000 are fetched. - var responses = new ArrayList(); - for (int i = 0; i <= 1000; i++) { - responses.add(toolsResponse(List.of(tool("t" + i)), "c" + i)); - } - var proxy = new FakeProxy(responses); - - assertThrows(IllegalStateException.class, proxy::listTools); - assertEquals(1000, proxy.requests.size()); - } - - @Test - void listToolsTreatsBlankCursorAsEndOfList() { - // A server that signals end-of-list with an empty cursor (instead of omitting it) must not - // trigger an extra round-trip or trip the repeated-cursor guard. - var proxy = new FakeProxy(List.of( - toolsResponse(List.of(tool("a"), tool("b")), ""))); - - var tools = proxy.listTools(); - - assertEquals(List.of("a", "b"), tools.stream().map(ToolInfo::getName).toList()); - assertEquals(1, proxy.requests.size()); - } - - @Test - void listToolsAbortsOnCyclingCursor() { - // A -> B -> A is a non-advancing cycle the consecutive-only check would miss; the - // seen-cursor guard must still abort it. - var proxy = new FakeProxy(List.of( - toolsResponse(List.of(tool("a")), "A"), - toolsResponse(List.of(tool("b")), "B"), - toolsResponse(List.of(tool("c")), "A"))); - - assertThrows(IllegalStateException.class, proxy::listTools); - } - - @Test - void listToolsThrowsOnErrorResponse() { - var proxy = new FakeProxy(List.of(errorResponse("boom"))); - - var ex = assertThrows(RuntimeException.class, proxy::listTools); - assertTrue(ex.getMessage().contains("boom")); - } - - @Test - void listToolsThrowsOnErrorOnLaterPage() { - var proxy = new FakeProxy(List.of( - toolsResponse(List.of(tool("a")), "c1"), - errorResponse("kaboom"))); - - assertThrows(RuntimeException.class, proxy::listTools); - assertEquals(2, proxy.requests.size()); - } - - @Test - void listToolsThrowsWhenResponseHasNeitherResultNorError() { - var proxy = new FakeProxy(List.of(emptyResponse())); - - var ex = assertThrows(RuntimeException.class, proxy::listTools); - assertTrue(ex.getMessage().contains("listing tools")); - } - - @Test - void listPromptsAbortsOnRepeatedCursor() { - var proxy = new FakeProxy(List.of( - promptsResponse(List.of(prompt("p1")), "SAME"), - promptsResponse(List.of(prompt("p2")), "SAME"))); - - assertThrows(IllegalStateException.class, proxy::listPrompts); - } - - @Test - void listToolsReturnsImmutableList() { - var proxy = new FakeProxy(List.of(toolsResponse(List.of(tool("a")), null))); - - var tools = proxy.listTools(); - - assertThrows(UnsupportedOperationException.class, () -> tools.add(tool("b"))); - } -} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java deleted file mode 100644 index 69006b214d..0000000000 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServiceTest.java +++ /dev/null @@ -1,243 +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 java.util.concurrent.TimeUnit.SECONDS; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotEquals; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.CyclicBarrier; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BooleanSupplier; -import org.junit.jupiter.api.Test; -import software.amazon.smithy.java.core.serde.document.Document; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; -import software.amazon.smithy.java.mcp.model.JsonRpcResponse; -import software.amazon.smithy.java.mcp.model.ListToolsResult; -import software.amazon.smithy.java.mcp.model.ToolInfo; - -class McpServiceTest { - - /** A proxy whose tool set can change and that records which thread its listTools() ran on. */ - private static final class FakeProxy extends McpServerProxy { - volatile List toolSet; - volatile String lastListToolsThread; - volatile CountDownLatch listToolsLatch = new CountDownLatch(1); - private final String name; - - FakeProxy(String name, List initial) { - this.name = name; - this.toolSet = initial; - } - - @Override - public List listTools() { - lastListToolsThread = Thread.currentThread().getName(); - listToolsLatch.countDown(); - return toolSet; - } - - @Override - protected CompletableFuture rpc(JsonRpcRequest request) { - return CompletableFuture.completedFuture(JsonRpcResponse.builder() - .jsonrpc("2.0") - .id(request.getId() == null ? Document.of(0) : request.getId()) - .result(Document.of(Map.of())) - .build()); - } - - @Override - protected void start() {} - - @Override - protected CompletableFuture shutdown() { - return CompletableFuture.completedFuture(null); - } - - @Override - public String name() { - return name; - } - - void fireListChanged() { - notify(JsonRpcRequest.builder() - .jsonrpc("2.0") - .method("notifications/tools/list_changed") - .build()); - } - } - - private static ToolInfo tool(String name) { - return ToolInfo.builder().name(name).build(); - } - - /** Builds a service with the fake proxy and drives initialize so its notification writer is wired. */ - private static McpService initializedService(FakeProxy proxy) { - var service = new McpService(Map.of(), - List.of(proxy), - "test", - "1.0", - (s, t) -> true, - null, - McpServerInterceptor.NOOP); - service.handleRequest( - JsonRpcRequest.builder() - .jsonrpc("2.0") - .id(Document.of(1)) - .method("initialize") - .params(Document.of(Map.of())) - .build(), - r -> {}, - ProtocolVersion.defaultVersion()); - return service; - } - - private static List listToolNames(McpService service) { - var resp = service.handleRequest( - JsonRpcRequest.builder().jsonrpc("2.0").id(Document.of(2)).method("tools/list").build(), - r -> {}, - ProtocolVersion.defaultVersion()); - return resp.getResult() - .asShape(ListToolsResult.builder()) - .getTools() - .stream() - .map(ToolInfo::getName) - .sorted() - .toList(); - } - - @Test - void listChangedRefreshRunsOffTheNotifyingThread() throws Exception { - // Regression for the reader-thread deadlock: a tools/list_changed refresh must NOT run - // listTools() on the thread that delivered the notification (on stdio that is the transport - // reader thread, which must stay free to read the tools/list response). - var proxy = new FakeProxy("fake", List.of(tool("a"))); - initializedService(proxy); - - // initialize() already called listTools() once on this thread; reset for the refresh. - proxy.lastListToolsThread = null; - proxy.listToolsLatch = new CountDownLatch(1); - - proxy.fireListChanged(); - - assertTrue(proxy.listToolsLatch.await(5, SECONDS), "refresh never ran"); - assertNotEquals(Thread.currentThread().getName(), - proxy.lastListToolsThread, - "refresh must not run on the notifying thread"); - assertTrue(proxy.lastListToolsThread != null && proxy.lastListToolsThread.startsWith("mcp-tools-refresh"), - "refresh should run on the dedicated executor thread, was: " + proxy.lastListToolsThread); - } - - @Test - void listChangedRefreshAddsNewToolsAndPrunesStaleOnes() throws Exception { - var proxy = new FakeProxy("fake", List.of(tool("old1"), tool("old2"))); - var service = initializedService(proxy); - assertEquals(List.of("old1", "old2"), listToolNames(service)); - - // Server's set changes: old1 kept, old2 gone, new1 added. - proxy.toolSet = List.of(tool("old1"), tool("new1")); - proxy.listToolsLatch = new CountDownLatch(1); - proxy.fireListChanged(); - assertTrue(proxy.listToolsLatch.await(5, SECONDS)); - - // The snapshot swap happens after listTools() returns, so poll for the expected state. - assertEventually(() -> List.of("new1", "old1").equals(listToolNames(service)), - "expected [new1, old1] but was " + listToolNames(service)); - } - - @Test - void concurrentRefreshAddProxyAndListDoNotLoseUpdatesOrThrow() throws Exception { - // Hammer the registry from multiple threads: repeated tools/list_changed refreshes on an - // existing proxy, dynamic addNewProxy calls, and concurrent tools/list reads. With - // copy-on-write under a single lock, reads must never throw and every added proxy's tool - // must be present at the end (no lost updates). - var refreshProxy = new FakeProxy("refresher", List.of(tool("r0"))); - var service = initializedService(refreshProxy); - - int adders = 4; - int refreshers = 4; - int readers = 4; - int iterations = 200; - var barrier = new CyclicBarrier(adders + refreshers + readers); - var error = new AtomicReference(); - var threads = new ArrayList(); - - for (int a = 0; a < adders; a++) { - final int id = a; - threads.add(new Thread(() -> { - try { - barrier.await(); - for (int i = 0; i < iterations; i++) { - var p = new FakeProxy("added-" + id + "-" + i, List.of(tool("added-" + id + "-" + i))); - service.addNewProxy(p, r -> {}); - } - } catch (Throwable t) { - error.compareAndSet(null, t); - } - })); - } - for (int r = 0; r < refreshers; r++) { - threads.add(new Thread(() -> { - try { - barrier.await(); - for (int i = 0; i < iterations; i++) { - refreshProxy.toolSet = List.of(tool("r" + i)); - service.refreshProxyTools(refreshProxy); - } - } catch (Throwable t) { - error.compareAndSet(null, t); - } - })); - } - for (int r = 0; r < readers; r++) { - threads.add(new Thread(() -> { - try { - barrier.await(); - for (int i = 0; i < iterations; i++) { - listToolNames(service); - } - } catch (Throwable t) { - error.compareAndSet(null, t); - } - })); - } - - threads.forEach(Thread::start); - for (var t : threads) { - t.join(30_000); - } - - if (error.get() != null) { - throw new AssertionError("concurrent access threw", error.get()); - } - - // Every proxy added by every adder thread must have its tool registered (no lost updates). - var finalNames = listToolNames(service); - for (int id = 0; id < adders; id++) { - var expected = "added-" + id + "-" + (iterations - 1); - assertTrue(finalNames.contains(expected), - "missing tool from a concurrently added proxy: " + expected); - } - } - - private static void assertEventually(BooleanSupplier condition, String message) - throws InterruptedException { - long deadline = System.nanoTime() + SECONDS.toNanos(5); - while (System.nanoTime() < deadline) { - if (condition.getAsBoolean()) { - return; - } - Thread.sleep(10); - } - throw new AssertionError(message); - } -} diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/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..0fe82f09d7 --- /dev/null +++ b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioMcpClientTest.java @@ -0,0 +1,62 @@ +/* + * 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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.math.BigInteger; +import java.time.Duration; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledOnOs; +import org.junit.jupiter.api.condition.OS; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.mcp.model.JsonRpcRequest; + +class 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"))); + } + + @Test + @EnabledOnOs({OS.LINUX, OS.MAC}) + void exchangeTimesOutWhenServerStaysSilent() { + var client = StdioMcpClient.builder() + .name("silent-server") + .command("sleep") + .arguments(List.of("30")) + .timeout(Duration.ofMillis(500)) + .build(); + client.start(); + try { + var error = assertThrows( + McpRemoteException.class, + () -> client.exchange(JsonRpcRequest.builder() + .jsonrpc("2.0") + .id(Document.of(1)) + .method("tools/list") + .build())); + assertTrue(error.getMessage().contains("Timed out")); + } finally { + client.close(); + } + } +} 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 80% 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 3d8c2df257..4a7dfd48ec 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 @@ -5,7 +5,6 @@ package software.amazon.smithy.java.mcp.server; -import static java.util.concurrent.TimeUnit.SECONDS; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -22,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; @@ -46,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) @@ -65,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) { @@ -77,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(); @@ -87,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) @@ -119,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) @@ -133,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(); @@ -155,9 +186,133 @@ public void initializeWithV2025_11_25ProtocolVersion() { assertTrue(outputSchema.get("properties").asStringMap().containsKey("outputStr")); } + @Test + public void doesNotAdvertiseStubUtilities() { + 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"); + assertNull(capabilities.getMember("completions")); + assertNull(capabilities.getMember("logging")); + assertTrue(capabilities.getMember("tools").getMember("listChanged").asBoolean()); + assertTrue(capabilities.getMember("prompts").getMember("listChanged").asBoolean()); + + 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")); + assertNull(discover.getMember("capabilities").getMember("completions")); + 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) @@ -171,7 +326,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(); @@ -189,7 +344,7 @@ public void noOutputSchemaWithUnsupportedProtocolVersion() { @Test public void validateToolsList() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -225,7 +380,7 @@ public void validateToolsList() { @Test public void validateNoIOOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -239,7 +394,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(); @@ -259,7 +414,7 @@ public void validateNoIOOperationTool() { @Test public void validateNoOutputOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -273,7 +428,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(); @@ -302,7 +457,7 @@ public void validateNoOutputOperationTool() { @Test public void validateNoInputOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -316,7 +471,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(); @@ -345,7 +500,7 @@ public void validateNoInputOperationTool() { @Test public void validateTestOperationTool() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -359,7 +514,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(); @@ -379,7 +534,7 @@ public void validateTestOperationTool() { @Test void readOnlyOperationHasReadOnlyHintAnnotation() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -411,7 +566,7 @@ void readOnlyOperationHasReadOnlyHintAnnotation() { @Test void idempotentOperationHasIdempotentHintAnnotation() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -443,7 +598,7 @@ void idempotentOperationHasIdempotentHintAnnotation() { @Test void plainOperationHasNoAnnotations() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -473,7 +628,7 @@ void plainOperationHasNoAnnotations() { @Test void annotationsStrippedForOldProtocolVersion() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -487,7 +642,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(); @@ -503,7 +658,7 @@ void annotationsStrippedForOldProtocolVersion() { @Test void testNumberAndStringIds() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -523,6 +678,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(); @@ -549,7 +718,7 @@ void testNumberAndStringIds() { @Test void testInvalidIds() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -592,7 +761,7 @@ void testInvalidIds() { @Test void testRequestsRequireIds() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -606,17 +775,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) @@ -639,7 +807,7 @@ public void readBeforeSerialization(InputHook hook) { server.start(); var bigDecimalValue = BigDecimal.valueOf(Integer.MAX_VALUE).add(BigDecimal.TEN); - var bigIntegerValue = BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.valueOf(100)); + var bigIntegerValue = BigInteger.valueOf(42); var blobValue = "Hello, World!"; var blobValueBase64 = Base64.getEncoder().encodeToString(blobValue.getBytes(StandardCharsets.UTF_8)); var nestedBigDecimalValue = new BigDecimal("123.456"); @@ -648,30 +816,32 @@ public void readBeforeSerialization(InputHook hook) { var nestedBlobValueBase64 = Base64.getEncoder().encodeToString(nestedBlobValue.getBytes(StandardCharsets.UTF_8)); + var nestedArguments = Document.of(Map.of( + "nestedBigDecimal", + Document.of(nestedBigDecimalValue), + "nestedBigInteger", + Document.of(nestedBigIntegerValue), + "nestedBlob", + Document.of(nestedBlobValueBase64), + "bigDecimalList", + Document.of(List.of( + Document.of(new BigDecimal("100.25")), + Document.of(new BigDecimal("200.75")))))); + var arguments = Document.of(Map.of( + "bigDecimalField", + Document.of(bigDecimalValue), + "bigIntegerField", + Document.of(42), + "blobField", + Document.of(blobValueBase64), + "nestedWithBigNumbers", + nestedArguments)); write("tools/call", - Document.of( - Map.of("name", - Document.of("TestOperation"), - "arguments", - Document.of(Map.of( - "bigDecimalField", - Document.of(bigDecimalValue.toString()), - "bigIntegerField", - Document.of(bigIntegerValue.toString()), - "blobField", - Document.of(blobValueBase64), - "nestedWithBigNumbers", - Document.of(Map.of( - "nestedBigDecimal", - Document.of(nestedBigDecimalValue.toString()), - "nestedBigInteger", - Document.of(nestedBigIntegerValue.toString()), - "nestedBlob", - Document.of(nestedBlobValueBase64), - "bigDecimalList", - Document.of(List.of( - Document.of("100.25"), - Document.of("200.75")))))))))); + Document.of(Map.of( + "name", + Document.of("TestOperation"), + "arguments", + arguments))); assertNotNull(read()); var inputDocument = capturedInput.get(); @@ -725,9 +895,69 @@ public void readBeforeSerialization(InputHook hook) { server.shutdown().join(); } + @Test + void invalidToolArgumentsReturnInvalidParams() { + server = StdioMcpServer.builder() + .input(input) + .output(output) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test#TestService")) + .proxyEndpoint("http://localhost") + .model(MODEL) + .build()) + .build(); + server.start(); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_11_25); + + write("tools/call", + Document.of(Map.of( + "name", + Document.of("TestOperation"), + "arguments", + Document.of(Map.of("list", Document.of("not-a-list")))))); + var response = read(); + + assertEquals(-32602, response.getError().getCode()); + assertTrue(response.getError().getMessage().contains("Invalid arguments for tool TestOperation")); + } + + @Test + void filteredToolsCannotBeListedOrCalled() { + server = StdioMcpServer.builder() + .input(input) + .output(output) + .toolFilter((serverId, toolName) -> !toolName.equals("NoIOOperation")) + .addService("test-mcp", + ProxyService.builder() + .service(ShapeId.from("smithy.test#TestService")) + .proxyEndpoint("http://localhost") + .model(MODEL) + .build()) + .build(); + server.start(); + initializeWithProtocolVersion(KnownProtocolVersion.V2025_11_25); + + write("tools/list", Document.of(Map.of())); + var tools = read().getResult().getMember("tools").asList(); + assertFalse(tools.stream() + .anyMatch(tool -> tool.getMember("name").asString().equals("NoIOOperation"))); + + write("tools/call", + Document.of(Map.of( + "name", + Document.of("NoIOOperation"), + "arguments", + Document.of(Map.of())))); + var response = read(); + + assertEquals(-32602, response.getError().getCode()); + assertEquals("No such tool: NoIOOperation", response.getError().getMessage()); + } + @Test void testNotificationsDoNotRequireRequestId() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -757,7 +987,7 @@ void testNotificationsDoNotRequireRequestId() { @Test void testUnknownMethodReturnsMethodNotFound() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -797,7 +1027,7 @@ void testUnknownMethodReturnsMethodNotFound() { @Test void testUnknownNotificationIsSilentlyDropped() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -828,7 +1058,7 @@ void testUnknownNotificationIsSilentlyDropped() { @Test void testPromptsList() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -872,7 +1102,7 @@ void testPromptsList() { @Test void testPromptsGetWithValidPrompt() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -906,7 +1136,7 @@ void testPromptsGetWithValidPrompt() { @Test void testPromptsGetWithDifferentCasing() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -977,7 +1207,7 @@ void testPromptsGetWithDifferentCasing() { @Test void testPromptsGetWithInvalidPrompt() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1017,7 +1247,7 @@ void testPromptsGetWithTemplateArguments() { .assemble() .unwrap(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1059,7 +1289,7 @@ void testPromptsGetWithMissingRequiredArguments() { .assemble() .unwrap(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1097,7 +1327,7 @@ void testApplyTemplateArgumentsEdgeCases() { .assemble() .unwrap(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1582,7 +1812,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) @@ -1596,7 +1826,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(); @@ -1637,7 +1867,7 @@ void testUnionSchemaGeneratesOneOfWithWrappedMembers() { @Test void testUnionWithOneOfTraitSchemaAlsoGeneratesOneOf() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -1651,7 +1881,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(); @@ -1677,13 +1907,13 @@ void testToolsListChangedNotificationInvalidatesCache() throws InterruptedExcept 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() @@ -1692,7 +1922,7 @@ void testToolsListChangedNotificationInvalidatesCache() throws InterruptedExcept .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"), @@ -1705,12 +1935,12 @@ void testToolsListChangedNotificationInvalidatesCache() throws InterruptedExcept .params(Document.of(Map.of())) .jsonrpc("2.0") .build(); - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(1, callCounter.get(), "Only initialize should have fetched from the proxy so far"); + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); + assertEquals(1, callCounter.get(), "First call should fetch from proxy"); - // Second tools/list - still just reads the registry, no proxy fetch. - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(1, callCounter.get(), "tools/list must not fetch from the proxy on its own"); + // Second tools/list - uses cache + service.execute(toolsRequest, ProtocolVersion.defaultVersion()); + assertEquals(1, callCounter.get(), "Second call should use cache"); // Send tools/list_changed notification. The refresh runs asynchronously off the notifying // thread (so the transport reader thread can't deadlock), so it fetches from the proxy a @@ -1726,17 +1956,19 @@ void testToolsListChangedNotificationInvalidatesCache() throws InterruptedExcept assertEquals(1, notifications.size()); assertEquals("notifications/tools/list_changed", notifications.get(0).getMethod()); - // The async refresh must eventually fetch from the proxy exactly once. - long deadline = System.nanoTime() + SECONDS.toNanos(5); - while (callCounter.get() < 2 && System.nanoTime() < deadline) { - Thread.sleep(10); - } - assertEquals(2, callCounter.get(), "list_changed should trigger exactly one refresh fetch"); + assertTimeoutPreemptively(Duration.ofSeconds(2), () -> { + while (callCounter.get() < 2) { + Thread.sleep(10); + } + }); - // Further tools/list calls just read the refreshed registry, no additional proxy fetch. - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - service.handleRequest(toolsRequest, r -> {}, ProtocolVersion.defaultVersion()); - assertEquals(2, callCounter.get(), "tools/list must not fetch from the proxy after the refresh"); + // 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.execute(toolsRequest, ProtocolVersion.defaultVersion()); + assertEquals(2, callCounter.get(), "Fourth call should use cache (not increment to 3)"); } @Test @@ -1744,13 +1976,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() @@ -1759,7 +1991,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"), @@ -1772,7 +2004,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 @@ -1787,11 +2019,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<>(); @@ -1800,34 +2032,33 @@ private static class CacheTestProxy extends McpServerProxy { } @Override - public List listTools() { + public McpPage listTools() { callCounter.incrementAndGet(); - return List.of( + return McpPage.last(List.of( ToolInfo.builder() .name("test-tool") .description("Test") .inputSchema(JsonObjectSchema.builder().build()) - .build()); + .build())); } @Override - public List listPrompts() { - return List.of(); + public McpPage listPrompts() { + return McpPage.last(List.of()); } @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() { @@ -1835,12 +2066,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() { @@ -1857,9 +2086,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) @@ -1869,16 +2098,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); @@ -1901,7 +2130,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) @@ -1911,18 +2140,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); @@ -1954,7 +2183,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) @@ -1964,18 +2193,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); @@ -2003,7 +2232,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) @@ -2013,11 +2242,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(); @@ -2045,7 +2274,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) @@ -2055,16 +2284,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); @@ -2084,10 +2313,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 @@ -2095,7 +2326,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) @@ -2105,12 +2336,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); @@ -2137,7 +2368,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) @@ -2147,26 +2378,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); } } @@ -2195,7 +2426,7 @@ void testContextPassesBetweenReadHooks() { Context.Key START_KEY = Context.key("start"); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2205,19 +2436,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); } }) @@ -2235,7 +2466,7 @@ public void readAfterExecution( @Test void testModifyBeforeExecutionRewritesRequest() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2245,15 +2476,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(); @@ -2268,7 +2494,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) @@ -2278,11 +2504,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(); @@ -2304,7 +2530,7 @@ public JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) { @Test void testModifyAfterExecutionTransformsResponse() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2314,21 +2540,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(); @@ -2342,7 +2566,7 @@ public JsonRpcResponse modifyAfterExecution( @Test void testModifyAfterToolCallTransformsResponse() { - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2352,19 +2576,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(); @@ -2389,7 +2611,7 @@ public JsonRpcResponse modifyAfterToolCall( void testReadBeforeExecutionThrowShortCircuits() { var afterExecutionError = new AtomicReference(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2399,16 +2621,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); @@ -2421,14 +2643,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) @@ -2438,24 +2661,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(); @@ -2474,7 +2693,7 @@ public JsonRpcResponse modifyAfterExecution( void testChainReadHooksInvokedInOrder() { var order = new ArrayList(); - server = McpServer.builder() + server = StdioMcpServer.builder() .name("smithy-mcp-server") .input(input) .output(output) @@ -2484,32 +2703,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"); @@ -2526,9 +2745,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) @@ -2538,26 +2758,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(); @@ -2574,9 +2793,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 @@ -2586,7 +2803,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) @@ -2596,22 +2813,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); @@ -2627,7 +2844,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 @@ -2637,7 +2855,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) @@ -2647,22 +2865,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); @@ -2685,6 +2903,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/StdioProxyTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioProxyTest.java deleted file mode 100644 index 34367f61de..0000000000 --- a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/StdioProxyTest.java +++ /dev/null @@ -1,48 +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 org.junit.jupiter.api.Assertions.assertInstanceOf; -import static org.junit.jupiter.api.Assertions.assertThrows; - -import java.time.Duration; -import java.util.List; -import java.util.concurrent.CompletionException; -import java.util.concurrent.TimeoutException; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.condition.EnabledOnOs; -import org.junit.jupiter.api.condition.OS; -import software.amazon.smithy.java.core.serde.document.Document; -import software.amazon.smithy.java.mcp.model.JsonRpcRequest; - -class StdioProxyTest { - - @Test - @EnabledOnOs({OS.LINUX, OS.MAC}) - void rpcTimesOutWhenServerStaysSilent() { - // `sleep` accepts the request on stdin but never writes a response, so the request future must - // fail via the per-request timeout rather than blocking the caller forever. - var proxy = StdioProxy.builder() - .name("silent-server") - .command("sleep") - .arguments(List.of("30")) - .timeout(Duration.ofMillis(500)) - .build(); - proxy.start(); - try { - var future = proxy.rpc(JsonRpcRequest.builder() - .jsonrpc("2.0") - .id(Document.of(1)) - .method("tools/list") - .build()); - - var ex = assertThrows(CompletionException.class, future::join); - assertInstanceOf(TimeoutException.class, ex.getCause()); - } finally { - proxy.shutdown().join(); - } - } -} 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..dacd8259d4 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,16 +34,38 @@ 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", "smithy.api#examples", + "smithy.ai#mcpHeader", "smithy.ai#prompts", "smithy.mcp#oneOf")) .build()) 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 +}