feat: Add support for native remote MCP servers (#943)

DivyaMaddipudi · web-flow · commit e22bb5c9768a · 2025-11-25T04:25:08.000Z
* feat: Add support for native remote MCP servers

* style: format code

* Updated HttpMcpProxy to use smithy-java HTTP client with dynamic header

* Refactor HttpMcpProxy to use Signer interface instead of headersSupplier, add Accept header, and return JsonRpcErrorResponse for non-2xx responses

* Clean up: rename url() to endpoint() and remove debug logs

* Refactor HttpMcpProxy: use Signer interface, improve error-handling, and optimize response processing

* Pass protocol version from initialize request to proxies and add truncation indicator for error responses

* Updated protocolVersion as volatile for thread-safe visibility
diff --git a/mcp/mcp-server/build.gradle.kts b/mcp/mcp-server/build.gradle.kts
@@ -16,9 +16,10 @@ dependencies {
     implementation(project(":codecs:json-codec", configuration = "shadow"))
     implementation(project(":mcp:mcp-schemas"))
     implementation(project(":smithy-ai-traits"))
+    implementation(project(":client:client-core"))
+    implementation(project(":client:client-http"))
     testRuntimeOnly(libs.smithy.aws.traits)
     testRuntimeOnly(project(":aws:client:aws-client-awsjson"))
-    testImplementation(project(":client:client-core"))
     testImplementation(project(":server:server-proxy"))
 }
 
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/HttpMcpProxy.java
@@ -0,0 +1,198 @@
+/*
+ * 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.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import software.amazon.smithy.java.auth.api.Signer;
+import software.amazon.smithy.java.client.core.ClientTransport;
+import software.amazon.smithy.java.client.http.HttpContext;
+import software.amazon.smithy.java.client.http.JavaHttpClientTransport;
+import software.amazon.smithy.java.context.Context;
+import software.amazon.smithy.java.http.api.HttpRequest;
+import software.amazon.smithy.java.http.api.HttpResponse;
+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.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 final ClientTransport<HttpRequest, HttpResponse> transport;
+    private final URI endpoint;
+    private final String name;
+    private final Signer<HttpRequest, ?> signer;
+    private final Duration timeout;
+
+    private HttpMcpProxy(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());
+        this.signer = builder.signer;
+        this.timeout = builder.timeout != null ? builder.timeout : Duration.ofSeconds(60);
+    }
+
+    private static String sanitizeName(String host) {
+        if (host == null) {
+            return "http-proxy-mcp";
+        }
+        return host.replaceAll("[^a-zA-Z0-9-]", "-");
+    }
+
+    public static final class Builder {
+        private String endpoint;
+        private String name;
+        private Signer<HttpRequest, ?> signer;
+        private ClientTransport<HttpRequest, HttpResponse> transport;
+        private Duration timeout;
+
+        public Builder endpoint(String endpoint) {
+            this.endpoint = endpoint;
+            return this;
+        }
+
+        public Builder name(String name) {
+            this.name = name;
+            return this;
+        }
+
+        public Builder signer(Signer<HttpRequest, ?> signer) {
+            this.signer = signer;
+            return this;
+        }
+
+        public Builder transport(ClientTransport<HttpRequest, HttpResponse> transport) {
+            this.transport = transport;
+            return this;
+        }
+
+        public Builder timeout(Duration timeout) {
+            this.timeout = timeout;
+            return this;
+        }
+
+        public HttpMcpProxy build() {
+            if (endpoint == null || endpoint.isEmpty()) {
+                throw new IllegalArgumentException("Endpoint must be provided");
+            }
+            return new HttpMcpProxy(this);
+        }
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    @Override
+    public CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
+        try {
+            byte[] body = JSON_CODEC.serializeToString(request).getBytes(StandardCharsets.UTF_8);
+            LOG.trace("Sending HTTP request to {}", endpoint);
+
+            String protocolVersionHeader = protocolVersion != null
+                    ? protocolVersion
+                    : ProtocolVersion.defaultVersion().identifier();
+
+            HttpRequest httpRequest = HttpRequest.builder()
+                    .uri(endpoint)
+                    .method("POST")
+                    .withAddedHeader("Content-Type", "application/json")
+                    .withAddedHeader("Accept", "application/json, text/event-stream")
+                    .withAddedHeader("MCP-Protocol-Version", protocolVersionHeader)
+                    .body(DataStream.ofBytes(body, "application/json"))
+                    .build();
+
+            Context context = Context.create();
+            context.put(HttpContext.HTTP_REQUEST_TIMEOUT, timeout);
+
+            if (signer != null) {
+                httpRequest = signer.sign(httpRequest, null, context);
+            }
+
+            HttpResponse response = transport.send(context, httpRequest);
+            LOG.trace("Received HTTP response with status: {}", response.statusCode());
+
+            if (response.statusCode() < 200 || response.statusCode() >= 300) {
+                return CompletableFuture.completedFuture(handleErrorResponse(response));
+            }
+
+            return CompletableFuture.completedFuture(JsonRpcResponse.builder()
+                    .deserialize(JSON_CODEC.createDeserializer(response.body().asByteBuffer()))
+                    .build());
+        } catch (Exception e) {
+            return CompletableFuture.failedFuture(e);
+        }
+    }
+
+    private JsonRpcResponse handleErrorResponse(HttpResponse response) {
+        long contentLength = response.body().contentLength();
+        String errorMessage = "HTTP " + response.statusCode();
+
+        if (contentLength > 0) {
+            String contentType = response.body().contentType();
+            byte[] bodyBytes = ByteBufferUtils.getBytes(response.body().asByteBuffer());
+
+            if ("application/json".equals(contentType)) {
+                try {
+                    return JsonRpcResponse.builder()
+                            .deserialize(JSON_CODEC.createDeserializer(bodyBytes))
+                            .build();
+                } catch (Exception e) {
+                    LOG.warn("Failed to deserialize JSON error response", e);
+                    return JsonRpcResponse.builder()
+                            .jsonrpc("2.0")
+                            .error(JsonRpcErrorResponse.builder()
+                                    .code(response.statusCode())
+                                    .message("HTTP " + response.statusCode() + ": Invalid JSON response")
+                                    .build())
+                            .build();
+                }
+            } else {
+                int length = Math.min(200, bodyBytes.length);
+                String errorBody = new String(bodyBytes, 0, length, StandardCharsets.UTF_8);
+                errorMessage = errorMessage + ": " + errorBody + (length != bodyBytes.length ? " (truncated)" : "");
+            }
+        }
+
+        return JsonRpcResponse.builder()
+                .jsonrpc("2.0")
+                .error(JsonRpcErrorResponse.builder()
+                        .code(response.statusCode())
+                        .message(errorMessage)
+                        .build())
+                .build();
+    }
+
+    @Override
+    public void start() {
+        // HTTP is connectionless, nothing to start
+        LOG.debug("HTTP MCP proxy started for endpoint: {}", endpoint);
+    }
+
+    @Override
+    public CompletableFuture<Void> shutdown() {
+        // HTTP client doesn't need explicit shutdown
+        LOG.debug("HTTP MCP proxy shutdown for endpoint: {}", endpoint);
+        return CompletableFuture.completedFuture(null);
+    }
+
+    @Override
+    public String name() {
+        return name;
+    }
+}
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
@@ -131,7 +131,7 @@ private void writeResponse(JsonRpcResponse response) {
     @Override
     public void start() {
         // Initialize proxies
-        mcpService.initializeProxies(this::writeResponse);
+        mcpService.startProxies();
 
         // Start the listener thread
         listener.start();
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
@@ -23,6 +23,7 @@ public abstract class McpServerProxy {
     private static final AtomicInteger ID_GENERATOR = new AtomicInteger(0);
 
     protected Consumer<JsonRpcResponse> notificationConsumer;
+    protected volatile String protocolVersion;
 
     public List<ToolInfo> listTools() {
         JsonRpcRequest request = JsonRpcRequest.builder()
@@ -43,13 +44,18 @@ public List<ToolInfo> listTools() {
         }).join();
     }
 
-    public void initialize(Consumer<JsonRpcResponse> notificationConsumer, JsonRpcRequest initializeRequest) {
+    public void initialize(
+            Consumer<JsonRpcResponse> notificationConsumer,
+            JsonRpcRequest initializeRequest,
+            String protocolVersion
+    ) {
 
         var result = Objects.requireNonNull(rpc(initializeRequest).join());
         if (result.getError() != null) {
             throw new RuntimeException("Error during initialization: " + result.getError().getMessage());
         }
         this.notificationConsumer = notificationConsumer;
+        this.protocolVersion = protocolVersion;
     }
 
     abstract CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request);
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
@@ -79,6 +79,7 @@ public final class McpService {
     private final Map<String, Service> services;
     private final AtomicReference<JsonRpcRequest> initializeRequest = new AtomicReference<>();
     private final ToolFilter toolFilter;
+    private volatile boolean proxiesInitialized = false;
     private final McpMetricsObserver metricsObserver;
 
     McpService(
@@ -166,6 +167,15 @@ private JsonRpcResponse handleInitialize(JsonRpcRequest req) {
         }
 
         this.initializeRequest.set(req);
+
+        // Initialize proxies lazily after we have a real initialize request
+        if (!proxiesInitialized) {
+            initializeProxies(response -> {
+                // Proxies are initialized, no additional response needed
+            });
+            proxiesInitialized = true;
+        }
+
         var maybeVersion = req.getParams().getMember("protocolVersion");
         String pv = null;
         if (maybeVersion != null) {
@@ -279,23 +289,47 @@ private JsonRpcResponse handleToolsCall(
     }
 
     /**
-     * Initializes proxies with the given response writer and initialize request.
+     * Starts proxies without initializing them.
      */
-    public void initializeProxies(Consumer<JsonRpcResponse> responseWriter) {
+    public void startProxies() {
         for (McpServerProxy proxy : proxies.values()) {
-            proxy.start();
-            proxy.initialize(responseWriter, initializeRequest.get());
+            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<JsonRpcResponse> responseWriter) {
+        JsonRpcRequest initRequest = initializeRequest.get();
+        if (initRequest == null) {
+            LOG.warn("Cannot initialize proxies: no initialize request received yet");
+            return;
+        }
 
-            // Get the tools from this proxy
+        String protocolVersion = null;
+        var maybeVersion = initRequest.getParams().getMember("protocolVersion");
+        if (maybeVersion != null) {
+            var version = ProtocolVersion.version(maybeVersion.asString());
+            if (!(version instanceof ProtocolVersion.UnknownVersion)) {
+                protocolVersion = version.identifier();
+            }
+        }
+
+        for (McpServerProxy proxy : proxies.values()) {
             try {
-                List<ToolInfo> proxyTools = proxy.listTools();
+                proxy.initialize(responseWriter, initRequest, protocolVersion);
 
-                // Add each tool to our tool map
+                List<ToolInfo> proxyTools = proxy.listTools();
                 for (var toolInfo : proxyTools) {
                     tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy));
                 }
             } catch (Exception e) {
-                LOG.error("Failed to fetch tools from proxy", e);
+                LOG.error("Failed to initialize proxy: " + proxy.name(), e);
             }
         }
     }
@@ -321,7 +355,6 @@ public void addNewService(String id, Service service) {
     public void addNewProxy(McpServerProxy mcpServerProxy, Consumer<JsonRpcResponse> responseWriter) {
         proxies.put(mcpServerProxy.name(), mcpServerProxy);
         mcpServerProxy.start();
-        mcpServerProxy.initialize(responseWriter, initializeRequest.get());
 
         try {
             List<ToolInfo> proxyTools = mcpServerProxy.listTools();
diff --git a/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java b/mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/HttpMcpProxyTest.java
