feat(mcp): Add AuthScheme and IdentityResolver support to HttpMcpProxy

Satyadivya Maddipudi · adwsingh · commit d488b60c98f6 · 2026-05-11T13:05:19.000-07:00
Add proper auth separation to HttpMcpProxy alongside the existing
signer() path. The new authScheme() + identityResolver() +
signerContext() builder API follows the canonical smithy-java auth
architecture (AuthScheme + IdentityResolver + Signer separation) used
by ClientPipeline. This enables direct use of SigV4AuthScheme and
SdkCredentialsResolver for AWS-authenticated MCP endpoints without
needing a custom adapter class.

The existing signer() path remains unchanged for backward compatibility.
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
@@ -10,7 +10,10 @@
 import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
 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;
 import software.amazon.smithy.java.client.core.ClientTransport;
+import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme;
 import software.amazon.smithy.java.client.http.HttpContext;
 import software.amazon.smithy.java.client.http.JavaHttpClientTransport;
 import software.amazon.smithy.java.context.Context;
@@ -42,6 +45,9 @@ public final class HttpMcpProxy extends McpServerProxy {
     private final URI endpoint;
     private final String name;
     private final Signer<HttpRequest, ?> signer;
+    private final AuthScheme<HttpRequest, ?> authScheme;
+    private final IdentityResolver<?> identityResolver;
+    private final Context signerContext;
     private final Duration timeout;
     private volatile String sessionId;
 
@@ -50,6 +56,9 @@ private HttpMcpProxy(Builder builder) {
         this.endpoint = URI.create(builder.endpoint);
         this.name = builder.name != null ? builder.name : sanitizeName(endpoint.getHost());
         this.signer = builder.signer;
+        this.authScheme = builder.authScheme;
+        this.identityResolver = builder.identityResolver;
+        this.signerContext = builder.signerContext != null ? builder.signerContext : Context.create();
         this.timeout = builder.timeout != null ? builder.timeout : Duration.ofMinutes(5);
     }
 
@@ -64,6 +73,9 @@ public static final class Builder {
         private String endpoint;
         private String name;
         private Signer<HttpRequest, ?> signer;
+        private AuthScheme<HttpRequest, ?> authScheme;
+        private IdentityResolver<?> identityResolver;
+        private Context signerContext;
         private ClientTransport<HttpRequest, HttpResponse> transport;
         private Duration timeout;
 
@@ -82,6 +94,21 @@ public Builder signer(Signer<HttpRequest, ?> signer) {
             return this;
         }
 
+        public Builder authScheme(AuthScheme<HttpRequest, ?> authScheme) {
+            this.authScheme = authScheme;
+            return this;
+        }
+
+        public Builder identityResolver(IdentityResolver<?> identityResolver) {
+            this.identityResolver = identityResolver;
+            return this;
+        }
+
+        public Builder signerContext(Context signerContext) {
+            this.signerContext = signerContext;
+            return this;
+        }
+
         public Builder transport(ClientTransport<HttpRequest, HttpResponse> transport) {
             this.transport = transport;
             return this;
@@ -96,6 +123,18 @@ public HttpMcpProxy build() {
             if (endpoint == null || endpoint.isEmpty()) {
                 throw new IllegalArgumentException("Endpoint must be provided");
             }
+            if (signer != null && authScheme != null) {
+                throw new IllegalArgumentException(
+                        "Cannot set both signer and authScheme; use one or the other");
+            }
+            if (authScheme != null && identityResolver == null) {
+                throw new IllegalArgumentException(
+                        "identityResolver must be provided when authScheme is set");
+            }
+            if (identityResolver != null && authScheme == null) {
+                throw new IllegalArgumentException(
+                        "authScheme must be provided when identityResolver is set");
+            }
             return new HttpMcpProxy(this);
         }
     }
@@ -137,7 +176,9 @@ public CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
             Context context = Context.create();
             context.put(HttpContext.HTTP_REQUEST_TIMEOUT, timeout);
 
-            if (signer != null) {
+            if (authScheme != null) {
+                httpRequest = signWithAuthScheme(httpRequest);
+            } else if (signer != null) {
                 httpRequest = signer.sign(httpRequest, null, context).signedRequest();
             }
 
@@ -178,6 +219,20 @@ public CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    private <I extends Identity> HttpRequest signWithAuthScheme(HttpRequest request) {
+        AuthScheme<HttpRequest, I> scheme = (AuthScheme<HttpRequest, I>) authScheme;
+        IdentityResolver<I> resolver = (IdentityResolver<I>) identityResolver;
+
+        Context signerProperties = scheme.getSignerProperties(signerContext);
+        Context identityProperties = scheme.getIdentityProperties(signerContext);
+        I identity = resolver.resolveIdentity(identityProperties).unwrap();
+
+        try (var schemeSigner = scheme.signer()) {
+            return schemeSigner.sign(request, identity, signerProperties).signedRequest();
+        }
+    }
+
     private JsonRpcResponse parseSseResponse(HttpResponse response, JsonRpcRequest request) {
         try {
             byte[] bodyBytes = ByteBufferUtils.getBytes(response.body().asByteBuffer());
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
@@ -21,10 +21,18 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.java.auth.api.SignResult;
+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;
+import software.amazon.smithy.java.auth.api.identity.IdentityResult;
+import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme;
+import software.amazon.smithy.java.context.Context;
 import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.java.http.api.HttpRequest;
 import software.amazon.smithy.java.json.JsonCodec;
 import software.amazon.smithy.java.mcp.model.JsonRpcRequest;
 import software.amazon.smithy.java.mcp.model.JsonRpcResponse;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeType;
 
 class HttpMcpProxyTest {
@@ -66,6 +74,110 @@ void testBuilderValidation() {
         assertThrows(IllegalArgumentException.class, () -> HttpMcpProxy.builder().endpoint("").build());
     }
 
+    @Test
+    void testBuilderRejectsSignerAndAuthSchemeTogether() {
+        assertThrows(IllegalArgumentException.class,
+                () -> HttpMcpProxy.builder()
+                        .endpoint(serverUrl)
+                        .signer((request, identity, context) -> new SignResult<>(request))
+                        .authScheme(new TestAuthScheme())
+                        .identityResolver(TestIdentityResolver.INSTANCE)
+                        .build());
+    }
+
+    @Test
+    void testBuilderRejectsAuthSchemeWithoutIdentityResolver() {
+        assertThrows(IllegalArgumentException.class,
+                () -> HttpMcpProxy.builder()
+                        .endpoint(serverUrl)
+                        .authScheme(new TestAuthScheme())
+                        .build());
+    }
+
+    @Test
+    void testBuilderRejectsIdentityResolverWithoutAuthScheme() {
+        assertThrows(IllegalArgumentException.class,
+                () -> HttpMcpProxy.builder()
+                        .endpoint(serverUrl)
+                        .identityResolver(TestIdentityResolver.INSTANCE)
+                        .build());
+    }
+
+    @Test
+    void testAuthSchemeSignsRequest() throws IOException {
+        String[] capturedHeader = {null};
+
+        mockServer.removeContext("/mcp");
+        mockServer.createContext("/mcp", exchange -> {
+            capturedHeader[0] = exchange.getRequestHeaders().getFirst("X-Test-Signed");
+            String response = "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":\"signed\"}";
+            exchange.getResponseHeaders().set("Content-Type", "application/json");
+            exchange.sendResponseHeaders(200, response.getBytes(StandardCharsets.UTF_8).length);
+            try (OutputStream os = exchange.getResponseBody()) {
+                os.write(response.getBytes(StandardCharsets.UTF_8));
+            }
+            exchange.close();
+        });
+
+        HttpMcpProxy authProxy = HttpMcpProxy.builder()
+                .endpoint(serverUrl)
+                .authScheme(new TestAuthScheme())
+                .identityResolver(TestIdentityResolver.INSTANCE)
+                .build();
+
+        JsonRpcRequest request = JsonRpcRequest.builder()
+                .method("test/method")
+                .id(Document.of(1))
+                .jsonrpc("2.0")
+                .build();
+
+        JsonRpcResponse response = authProxy.rpc(request).join();
+
+        assertNotNull(response);
+        assertEquals("signed", response.getResult().asString());
+        assertEquals("test-token", capturedHeader[0]);
+        authProxy.shutdown().join();
+    }
+
+    @Test
+    void testAuthSchemeReceivesSignerContext() throws IOException {
+        String[] capturedRegion = {null};
+
+        mockServer.removeContext("/mcp");
+        mockServer.createContext("/mcp", exchange -> {
+            capturedRegion[0] = exchange.getRequestHeaders().getFirst("X-Region");
+            String response = "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":\"ok\"}";
+            exchange.getResponseHeaders().set("Content-Type", "application/json");
+            exchange.sendResponseHeaders(200, response.getBytes(StandardCharsets.UTF_8).length);
+            try (OutputStream os = exchange.getResponseBody()) {
+                os.write(response.getBytes(StandardCharsets.UTF_8));
+            }
+            exchange.close();
+        });
+
+        Context signerCtx = Context.create();
+        signerCtx.put(TestAuthScheme.REGION_KEY, "us-west-2");
+
+        HttpMcpProxy authProxy = HttpMcpProxy.builder()
+                .endpoint(serverUrl)
+                .authScheme(new TestAuthScheme())
+                .identityResolver(TestIdentityResolver.INSTANCE)
+                .signerContext(signerCtx)
+                .build();
+
+        JsonRpcRequest request = JsonRpcRequest.builder()
+                .method("test/method")
+                .id(Document.of(1))
+                .jsonrpc("2.0")
+                .build();
+
+        JsonRpcResponse response = authProxy.rpc(request).join();
+
+        assertNotNull(response);
+        assertEquals("us-west-2", capturedRegion[0]);
+        authProxy.shutdown().join();
+    }
+
     @Test
     void testBuilderWithCustomName() {
         HttpMcpProxy customProxy = HttpMcpProxy.builder()
@@ -555,4 +667,64 @@ public void handle(HttpExchange exchange) throws IOException {
             }
         }
     }
+
+    private record TestIdentity(String token) implements Identity {}
+
+    private static final class TestIdentityResolver implements IdentityResolver<TestIdentity> {
+        static final TestIdentityResolver INSTANCE = new TestIdentityResolver();
+
+        @Override
+        public IdentityResult<TestIdentity> resolveIdentity(Context requestProperties) {
+            return IdentityResult.of(new TestIdentity("test-token"));
+        }
+
+        @Override
+        public Class<TestIdentity> identityType() {
+            return TestIdentity.class;
+        }
+    }
+
+    private static final class TestAuthScheme implements AuthScheme<HttpRequest, TestIdentity> {
+        static final Context.Key<String> REGION_KEY = Context.key("test-region");
+
+        @Override
+        public ShapeId schemeId() {
+            return ShapeId.from("smithy.test#testAuth");
+        }
+
+        @Override
+        public Class<HttpRequest> requestClass() {
+            return HttpRequest.class;
+        }
+
+        @Override
+        public Class<TestIdentity> identityClass() {
+            return TestIdentity.class;
+        }
+
+        @Override
+        public Context getSignerProperties(Context context) {
+            var ctx = Context.create();
+            var region = context.get(REGION_KEY);
+            if (region != null) {
+                ctx.put(REGION_KEY, region);
+            }
+            return ctx;
+        }
+
+        @Override
+        public Signer<HttpRequest, TestIdentity> signer() {
+            return (request, identity, properties) -> {
+                var r = request.toModifiable();
+                var h = r.headers().toModifiable();
+                h.setHeader("X-Test-Signed", identity.token());
+                var region = properties.get(REGION_KEY);
+                if (region != null) {
+                    h.setHeader("X-Region", region);
+                }
+                r.setHeaders(h);
+                return new SignResult<>(r);
+            };
+        }
+    }
 }
