Don't make it a hard dependency that initialize needs to be sent

Author: adwsingh

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpProxy.java

@@ -104,9 +104,7 @@ public CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
             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();
+            String protocolVersionHeader = getProtocolVersion().identifier();
 
             HttpRequest httpRequest = HttpRequest.builder()
                     .uri(endpoint)

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java

@@ -9,6 +9,7 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import software.amazon.smithy.java.core.schema.SerializableStruct;
 import software.amazon.smithy.java.core.schema.ShapeBuilder;
@@ -22,8 +23,9 @@ public abstract class McpServerProxy {
 
     private static final AtomicInteger ID_GENERATOR = new AtomicInteger(0);
 
-    protected Consumer<JsonRpcResponse> notificationConsumer;
-    protected volatile String protocolVersion;
+    private final AtomicReference<Consumer<JsonRpcResponse>> notificationConsumer = new AtomicReference<>();
+    private final AtomicReference<ProtocolVersion> protocolVersion =
+            new AtomicReference<>(ProtocolVersion.defaultVersion());
 
     public List<ToolInfo> listTools() {
         JsonRpcRequest request = JsonRpcRequest.builder()
@@ -47,15 +49,19 @@ public List<ToolInfo> listTools() {
     public void initialize(
             Consumer<JsonRpcResponse> notificationConsumer,
             JsonRpcRequest initializeRequest,
-            String protocolVersion
+            ProtocolVersion 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;
+        this.notificationConsumer.set(notificationConsumer);
+        this.protocolVersion.set(protocolVersion);
+    }
+
+    protected final ProtocolVersion getProtocolVersion() {
+        return protocolVersion.get();
     }
 
     abstract CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request);
@@ -85,7 +91,10 @@ protected Document generateRequestId() {
     }
 
     protected void notify(JsonRpcResponse response) {
-        notificationConsumer.accept(response);
+        var nc = notificationConsumer.get();
+        if (nc != null) {
+            nc.accept(response);
+        }
     }
 
     public abstract String name();

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java

@@ -79,7 +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 AtomicReference<Boolean> proxiesInitialized = new AtomicReference<>(false);
     private final McpMetricsObserver metricsObserver;
 
     McpService(
@@ -118,13 +118,19 @@ public JsonRpcResponse handleRequest(
     ) {
         try {
             validate(req);
-            return switch (req.getMethod()) {
+            var method = req.getMethod();
+            return switch (method) {
                 case "initialize" -> handleInitialize(req);
-                case "prompts/list" -> handlePromptsList(req);
-                case "prompts/get" -> handlePromptsGet(req);
-                case "tools/list" -> handleToolsList(req, protocolVersion);
-                case "tools/call" -> handleToolsCall(req, asyncResponseCallback, protocolVersion);
-                default -> null; // Notifications or unknown methods
+                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" -> handleToolsCall(req, asyncResponseCallback, protocolVersion);
+                        default -> null; // Notifications or unknown methods
+                    };
+                }
             };
         } catch (Exception e) {
             return createErrorResponse(req, e);
@@ -166,15 +172,9 @@ private JsonRpcResponse handleInitialize(JsonRpcRequest req) {
                     clientTitle);
         }
 
-        this.initializeRequest.set(req);
+        this.initializeRequest.compareAndSet(null, req);
 
-        // Initialize proxies lazily after we have a real initialize request
-        if (!proxiesInitialized) {
-            initializeProxies(response -> {
-                // Proxies are initialized, no additional response needed
-            });
-            proxiesInitialized = true;
-        }
+        initializeProxies(rpcResponse -> {});
 
         var maybeVersion = req.getParams().getMember("protocolVersion");
         String pv = null;
@@ -305,31 +305,32 @@ public void startProxies() {
      * 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;
-        }
-
-        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();
+        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()) {
-            try {
-                proxy.initialize(responseWriter, initRequest, protocolVersion);
+            for (McpServerProxy proxy : proxies.values()) {
+                try {
+                    if (initRequest != null) {
+                        proxy.initialize(responseWriter, initRequest, protocolVersion);
+                    }
 
-                List<ToolInfo> proxyTools = proxy.listTools();
-                for (var toolInfo : proxyTools) {
-                    tools.put(toolInfo.getName(), new Tool(toolInfo, proxy.name(), proxy));
+                    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 initialize proxy: " + proxy.name(), e);
                 }
-            } catch (Exception e) {
-                LOG.error("Failed to initialize proxy: " + proxy.name(), e);
             }
         }
     }