From d19bfebab9914f4cc2f5480fd53018e238f9b1a5 Mon Sep 17 00:00:00 2001
From: Minghui Wang <143184745+wmh123wmh1@users.noreply.github.com>
Date: Mon, 31 Aug 2026 18:00:32 +0800
Subject: [PATCH] mcp: preserve single-request batch responses (#1214)

## What does this PR do?

Preserve JSON-RPC batch response arrays when a batch contains a single
request.

## Why is this needed?

A single-request legacy batch must be returned as a one-element JSON
array, but the JSON response path currently unwraps it into a single
JSON object.

## Changes

- Preserve the original batch state on the logical stream
- Keep single-request batch responses as arrays
- Add a regression test for legacy initialization batches

## Testing

- `go test ./mcp -count=1`
- `go test ./... -count=1`
- `git diff --check`

Fixes #1211
---
 mcp/streamable.go      |  7 ++++++-
 mcp/streamable_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/mcp/streamable.go b/mcp/streamable.go
index b4a2a32..c8cd28b 100644
--- a/mcp/streamable.go
+++ b/mcp/streamable.go
@@ -988,6 +988,10 @@ type stream struct {
 	// Note: if we remove support for batching, this could just be a bool.
 	pendingJSONMessages []json.RawMessage
 
+	// isBatch reports whether the request that opened the stream was a JSON-RPC
+	// batch. Batch responses must remain arrays even when they contain one item.
+	isBatch bool
+
 	// w is the HTTP response writer for this stream. A non-nil w indicates
 	// that the stream is claimed by an HTTP request (the hanging POST or GET);
 	// it is set to nil when the request completes.
@@ -1138,7 +1142,7 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I
 		if done {
 			// Flush all pending messages as JSON response.
 			var toWrite []byte
-			if len(s.pendingJSONMessages) == 1 {
+			if len(s.pendingJSONMessages) == 1 && !s.isBatch {
 				toWrite = s.pendingJSONMessages[0]
 			} else {
 				toWrite, err = json.Marshal(s.pendingJSONMessages)
@@ -1691,6 +1695,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 		return
 	}
 	stream.isListen = isSubscriptionsListen
+	stream.isBatch = isBatch
 
 	// subscriptions/listen is inherently a long-lived SSE endpoint (SEP-2575):
 	// it has no synchronous result, the response stream stays open until the
diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go
index ff96f61..58e5ca7 100644
--- a/mcp/streamable_test.go
+++ b/mcp/streamable_test.go
@@ -4040,6 +4040,52 @@ func TestStreamableServerRejectsDuplicateInFlightRequestID(t *testing.T) {
 	}
 }
 
+// TestStreamableServerPreservesSingleRequestBatchResponse verifies that a
+// single-request JSON-RPC batch is returned as a one-element JSON array.
+func TestStreamableServerPreservesSingleRequestBatchResponse(t *testing.T) {
+	server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
+	handler := NewStreamableHTTPHandler(
+		func(*http.Request) *Server { return server },
+		&StreamableHTTPOptions{JSONResponse: true},
+	)
+	defer handler.closeAll()
+
+	httpServer := httptest.NewServer(handler)
+	defer httpServer.Close()
+
+	body := []byte(`[{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0"}}}]`)
+	req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json, text/event-stream")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	data, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("status = %d, want %d; body = %s", resp.StatusCode, http.StatusOK, data)
+	}
+
+	var messages []json.RawMessage
+	if err := json.Unmarshal(data, &messages); err != nil {
+		t.Fatalf("response is not a JSON array: %v; body = %s", err, data)
+	}
+	if len(messages) != 1 {
+		t.Fatalf("response contains %d messages, want 1; body = %s", len(messages), data)
+	}
+	if _, err := jsonrpc2.DecodeMessage(messages[0]); err != nil {
+		t.Fatalf("response item is not a JSON-RPC message: %v", err)
+	}
+}
+
 // TestStreamableMaxRequestBodyBytes verifies that the streamable HTTP handler
 // enforces StreamableHTTPOptions.MaxRequestBodyBytes on incoming request
 // bodies. The limit must apply uniformly regardless of transfer encoding: