Make McpService take a protocol version in every request. Default to 2025-03-26 protocol version

Author: adwsingh

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

@@ -37,6 +37,7 @@ public final class McpServer implements Server {
     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;
@@ -69,7 +70,18 @@ private void listen() {
     }
 
     private void handleRequest(JsonRpcRequest req) {
-        var response = mcpService.handleRequest(req, this::writeResponse);
+        // 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::writeResponse, protocolVersion);
         if (response != null) {
             writeResponse(response);
         }

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

@@ -78,7 +78,6 @@ public final class McpService {
     private final Map<String, Service> services;
     private final AtomicReference<JsonRpcRequest> initializeRequest = new AtomicReference<>();
     private final ToolFilter toolFilter;
-    private volatile ProtocolVersion protocolVersion;
 
     McpService(Map<String, Service> services, List<McpServerProxy> proxyList, String name, ToolFilter toolFilter) {
         this.services = services;
@@ -97,17 +96,22 @@ public final class McpService {
      *
      * @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 operations
      */
-    public JsonRpcResponse handleRequest(JsonRpcRequest req, Consumer<JsonRpcResponse> asyncResponseCallback) {
+    public JsonRpcResponse handleRequest(
+            JsonRpcRequest req,
+            Consumer<JsonRpcResponse> asyncResponseCallback,
+            ProtocolVersion protocolVersion
+    ) {
         try {
             validate(req);
             return switch (req.getMethod()) {
                 case "initialize" -> handleInitialize(req);
                 case "prompts/list" -> handlePromptsList(req);
                 case "prompts/get" -> handlePromptsGet(req);
-                case "tools/list" -> handleToolsList(req);
-                case "tools/call" -> handleToolsCall(req, asyncResponseCallback);
+                case "tools/list" -> handleToolsList(req, protocolVersion);
+                case "tools/call" -> handleToolsCall(req, asyncResponseCallback, protocolVersion);
                 default -> null; // Notifications or unknown methods
             };
         } catch (Exception e) {
@@ -120,7 +124,7 @@ private JsonRpcResponse handleInitialize(JsonRpcRequest req) {
         var maybeVersion = req.getParams().getMember("protocolVersion");
         String pv = null;
         if (maybeVersion != null) {
-            protocolVersion = ProtocolVersion.version(maybeVersion.asString());
+            var protocolVersion = ProtocolVersion.version(maybeVersion.asString());
             if (!(protocolVersion instanceof ProtocolVersion.UnknownVersion)) {
                 pv = protocolVersion.identifier();
             }
@@ -166,8 +170,8 @@ private JsonRpcResponse handlePromptsGet(JsonRpcRequest req) {
         return createSuccessResponse(req.getId(), result);
     }
 
-    private JsonRpcResponse handleToolsList(JsonRpcRequest req) {
-        var supportsOutputSchema = supportsOutputSchema();
+    private JsonRpcResponse handleToolsList(JsonRpcRequest req, ProtocolVersion protocolVersion) {
+        var supportsOutputSchema = supportsOutputSchema(protocolVersion);
         var result = ListToolsResult.builder()
                 .tools(tools.values()
                         .stream()
@@ -178,7 +182,11 @@ private JsonRpcResponse handleToolsList(JsonRpcRequest req) {
         return createSuccessResponse(req.getId(), result);
     }
 
-    private JsonRpcResponse handleToolsCall(JsonRpcRequest req, Consumer<JsonRpcResponse> asyncResponseCallback) {
+    private JsonRpcResponse handleToolsCall(
+            JsonRpcRequest req,
+            Consumer<JsonRpcResponse> asyncResponseCallback,
+            ProtocolVersion protocolVersion
+    ) {
         var operationName = req.getParams().getMember("name").asString();
         var tool = tools.get(operationName);
 
@@ -209,7 +217,7 @@ private JsonRpcResponse handleToolsCall(JsonRpcRequest req, Consumer<JsonRpcResp
             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);
+            var result = formatStructuredContent(tool, (SerializableShape) output, protocolVersion);
             return createSuccessResponse(req.getId(), result);
         }
     }
@@ -283,17 +291,21 @@ public Map<String, McpServerProxy> getProxies() {
         return proxies;
     }
 
-    private boolean supportsOutputSchema() {
+    private boolean supportsOutputSchema(ProtocolVersion protocolVersion) {
         return protocolVersion != null && protocolVersion.compareTo(ProtocolVersion.v2025_06_18.INSTANCE) >= 0;
     }
 
-    private CallToolResult formatStructuredContent(Tool tool, SerializableShape output) {
+    private CallToolResult formatStructuredContent(
+            Tool tool,
+            SerializableShape output,
+            ProtocolVersion protocolVersion
+    ) {
         var result = CallToolResult.builder()
                 .content(List.of(TextContent.builder()
                         .text(CODEC.serializeToString(output))
                         .build()));
 
-        if (supportsOutputSchema()) {
+        if (supportsOutputSchema(protocolVersion)) {
             var outputSchema = tool.toolInfo().getOutputSchema();
             if (outputSchema != null) {
                 result.structuredContent(Document.of(output));

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

@@ -62,10 +62,15 @@ public final int compareTo(ProtocolVersion o) {
 
     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;
             default -> new UnknownVersion(identifier);
         };
     }
+
+    public static ProtocolVersion defaultVersion() {
+        return v2025_03_26.INSTANCE;
+    }
 }