From 809fd4379fb100c9c6f3bf9a4e964ced4477e287 Mon Sep 17 00:00:00 2001
From: ankitm123 <ankitmohapatra123@gmail.com>
Date: Thu, 11 Sep 2025 10:05:43 +0200
Subject: [PATCH] mcp: remove json-rpc batching for more recent protocol
 versions

JSON-RPC batching support is removed for protocol versions greater than or
equal to protocolVersion20250618

Fixes #21
---
 mcp/streamable.go      | 19 ++++++++++++-------
 mcp/streamable_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/mcp/streamable.go b/mcp/streamable.go
index 1e35ca5b..bfaccae4 100644
--- a/mcp/streamable.go
+++ b/mcp/streamable.go
@@ -614,17 +614,22 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 		http.Error(w, "POST requires a non-empty body", http.StatusBadRequest)
 		return
 	}
-	// TODO(#21): if the negotiated protocol version is 2025-06-18 or later,
-	// we should not allow batching here.
-	//
-	// This also requires access to the negotiated version, which would either be
-	// set by the MCP-Protocol-Version header, or would require peeking into the
-	// session.
-	incoming, _, err := readBatch(body)
+	incoming, isBatch, err := readBatch(body)
 	if err != nil {
 		http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
 		return
 	}
+
+	protocolVersion := req.Header.Get(protocolVersionHeader)
+	if protocolVersion == "" {
+		protocolVersion = protocolVersion20250326
+	}
+
+	if isBatch && protocolVersion >= protocolVersion20250618 {
+		http.Error(w, fmt.Sprintf("JSON-RPC batching is not supported in %s and later (request version: %s)", protocolVersion20250618, protocolVersion), http.StatusBadRequest)
+		return
+	}
+
 	requests := make(map[jsonrpc.ID]struct{})
 	tokenInfo := auth.TokenInfoFromContext(req.Context())
 	isInitialize := false
diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go
index f731a083..8817784d 100644
--- a/mcp/streamable_test.go
+++ b/mcp/streamable_test.go
@@ -647,6 +647,46 @@ func TestStreamableServerTransport(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "batch rejected on 2025-06-18",
+			requests: []streamableRequest{
+				initialize,
+				initialized,
+				{
+					method: "POST",
+					// Explicitly set the protocol version header
+					headers: http.Header{"MCP-Protocol-Version": {"2025-06-18"}},
+					// Two messages => batch. Expect reject.
+					messages: []jsonrpc.Message{
+						req(101, "tools/call", &CallToolParams{Name: "tool"}),
+						req(102, "tools/call", &CallToolParams{Name: "tool"}),
+					},
+					wantStatusCode:     http.StatusBadRequest,
+					wantBodyContaining: "batch",
+				},
+			},
+		},
+		{
+			name: "batch accepted on 2025-03-26",
+			requests: []streamableRequest{
+				initialize,
+				initialized,
+				{
+					method:  "POST",
+					headers: http.Header{"MCP-Protocol-Version": {"2025-03-26"}},
+					// Two messages => batch. Expect OK with two responses in order.
+					messages: []jsonrpc.Message{
+						req(201, "tools/call", &CallToolParams{Name: "tool"}),
+						req(202, "tools/call", &CallToolParams{Name: "tool"}),
+					},
+					wantStatusCode: http.StatusOK,
+					wantMessages: []jsonrpc.Message{
+						resp(201, &CallToolResult{Content: []Content{}}, nil),
+						resp(202, &CallToolResult{Content: []Content{}}, nil),
+					},
+				},
+			},
+		},
 		{
 			name: "tool notification",
 			tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {