mcp: include decoded batch metadata in Streamable HTTP request summaries

mcp/streamable.go

@@ -69,22 +69,33 @@ type sessionInfo struct {
 	timer   *time.Timer
 }
 
-// StreamableHTTPRequestSummary contains redacted metadata about a single
-// JSON-RPC message decoded from a streamable HTTP POST body.
+// StreamableHTTPRequestSummary contains redacted metadata about JSON-RPC
+// messages decoded from a streamable HTTP POST body.
 type StreamableHTTPRequestSummary struct {
 	// Method is the method of a decoded JSON-RPC request. It is empty for a
-	// response. The method has not yet been validated and may contain an
-	// arbitrary attacker-controlled string.
+	// response or batch. The method has not yet been validated and may contain
+	// an arbitrary attacker-controlled string.
 	Method string
 
-	// RequestID is valid only for a JSON-RPC call. IDs use the same coercion
-	// rules as [jsonrpc.DecodeMessage].
+	// RequestID is valid only for a single JSON-RPC call. IDs use the same
+	// coercion rules as [jsonrpc.DecodeMessage].
 	RequestID jsonrpc.ID
 
-	// IsNotification reports whether the message is a JSON-RPC notification.
+	// BatchCount is the number of messages in a JSON-RPC batch. It is zero for
+	// a single message, including a single-message POST that is not an array.
+	BatchCount int
+
+	// Methods contains the lexicographically sorted unique request methods in a
+	// batch. Responses have no method and are excluded. Methods are unvalidated
+	// attacker-controlled strings. It is nil for a single message.
+	Methods []string
+
+	// IsNotification reports whether a single message is a JSON-RPC notification
+	// or every message in a batch is a notification.
 	IsNotification bool
 
-	// IsResponse reports whether the message is a JSON-RPC response.
+	// IsResponse reports whether a single message is a JSON-RPC response. It is
+	// false for a batch, including one containing only responses.
 	IsResponse bool
 }
 
@@ -237,10 +248,10 @@ type StreamableHTTPOptions struct {
 	// the allowsessionsinstateless compatibility path) are unaffected.
 	PropagateRequestCancellation bool
 
-	// OnRequestSummary, when non-nil, observes redacted metadata for a single
-	// JSON-RPC message decoded from a streamable HTTP POST body. It is not called
-	// for JSON-RPC batches. The callback receives the HTTP request's context and
-	// runs synchronously before validation and dispatch of the decoded message.
+	// OnRequestSummary, when non-nil, observes redacted metadata for a JSON-RPC
+	// message or batch decoded from a streamable HTTP POST body. The callback
+	// receives the HTTP request's context and runs synchronously before protocol
+	// version validation and dispatch of the decoded messages.
 	// It may be called concurrently for different requests and should return
 	// promptly; in particular, it must not wait for processing of the same
 	// request. Panics are not recovered. Only the summary is redacted; the context
@@ -1612,8 +1623,12 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 		http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
 		return
 	}
-	if c.onRequestSummary != nil && !isBatch && len(incoming) == 1 {
-		c.onRequestSummary(req.Context(), summarizeStreamableHTTPRequest(incoming[0]))
+	if c.onRequestSummary != nil {
+		if isBatch {
+			c.onRequestSummary(req.Context(), summarizeStreamableHTTPRequestBatch(incoming))
+		} else if len(incoming) == 1 {
+			c.onRequestSummary(req.Context(), summarizeStreamableHTTPRequest(incoming[0]))
+		}
 	}
 
 	protocolVersion := protocolVersionFromContext(req.Context())
@@ -1967,6 +1982,27 @@ func summarizeStreamableHTTPRequest(msg jsonrpc.Message) StreamableHTTPRequestSu
 	return summary
 }
 
+func summarizeStreamableHTTPRequestBatch(messages []jsonrpc.Message) StreamableHTTPRequestSummary {
+	summary := StreamableHTTPRequestSummary{BatchCount: len(messages), IsNotification: true}
+	seen := make(map[string]bool)
+	for _, message := range messages {
+		request, ok := message.(*jsonrpc.Request)
+		if !ok {
+			summary.IsNotification = false
+			continue
+		}
+		if request.IsCall() {
+			summary.IsNotification = false
+		}
+		if !seen[request.Method] {
+			seen[request.Method] = true
+			summary.Methods = append(summary.Methods, request.Method)
+		}
+	}
+	slices.Sort(summary.Methods)
+	return summary
+}
+
 // Event IDs: encode both the logical connection ID and the index, as
 // <streamID>_<idx>, to be consistent with the typescript implementation.
 

mcp/streamable_test.go

@@ -4302,6 +4302,60 @@ func TestSummarizeStreamableHTTPRequest(t *testing.T) {
 	}
 }
 
+func TestSummarizeStreamableHTTPRequestBatch(t *testing.T) {
+	tests := []struct {
+		name             string
+		body             string
+		wantMethods      []string
+		wantCount        int
+		wantNotification bool
+	}{
+		{
+			name:        "mixed calls and notifications",
+			body:        `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","id":2,"method":"ping"},{"jsonrpc":"2.0","id":3,"method":"tools/list"}]`,
+			wantMethods: []string{"notifications/initialized", "ping", "tools/list"},
+			wantCount:   4,
+		},
+		{
+			name:             "notification-only batch",
+			body:             `[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"notifications/initialized"}]`,
+			wantMethods:      []string{"notifications/initialized"},
+			wantCount:        2,
+			wantNotification: true,
+		},
+		{
+			name:        "single-item batch",
+			body:        `[{"jsonrpc":"2.0","id":"one","method":"ping"}]`,
+			wantMethods: []string{"ping"},
+			wantCount:   1,
+		},
+		{
+			name:        "response prevents notification-only classification",
+			body:        `[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","id":1,"result":{}}]`,
+			wantMethods: []string{"notifications/initialized"},
+			wantCount:   2,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			messages, isBatch, err := readBatch([]byte(test.body))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !isBatch {
+				t.Fatal("input was not decoded as a batch")
+			}
+			got := summarizeStreamableHTTPRequestBatch(messages)
+			if got.Method != "" || got.RequestID.IsValid() || got.IsResponse {
+				t.Errorf("single-message fields populated for batch: %+v", got)
+			}
+			if got.BatchCount != test.wantCount || got.IsNotification != test.wantNotification || !slices.Equal(got.Methods, test.wantMethods) {
+				t.Errorf("summary = %+v, want count %d, methods %q, notification %t", got, test.wantCount, test.wantMethods, test.wantNotification)
+			}
+		})
+	}
+}
+
 func TestStreamableHTTPRequestSummaryContext(t *testing.T) {
 	type contextKey struct{}
 	const contextValue = "middleware-value"
@@ -4388,12 +4442,12 @@ func TestStreamableHTTPRequestSummaryRejections(t *testing.T) {
 		}
 	})
 
-	t.Run("batch is not observed", func(t *testing.T) {
-		var calls atomic.Int64
+	t.Run("decoded batch rejected before dispatch is observed", func(t *testing.T) {
+		observed := make(chan StreamableHTTPRequestSummary, 1)
 		handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{
 			Stateless: true,
-			OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) {
-				calls.Add(1)
+			OnRequestSummary: func(_ context.Context, summary StreamableHTTPRequestSummary) {
+				observed <- summary
 			},
 		})
 		req := newRequest(`[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"notifications/initialized"}]`)
@@ -4403,8 +4457,33 @@ func TestStreamableHTTPRequestSummaryRejections(t *testing.T) {
 		if rec.Code != http.StatusBadRequest {
 			t.Errorf("status = %d, want %d", rec.Code, http.StatusBadRequest)
 		}
-		if got := calls.Load(); got != 0 {
-			t.Errorf("callback calls = %d, want 0", got)
+		select {
+		case got := <-observed:
+			if got.BatchCount != 2 || !got.IsNotification || !slices.Equal(got.Methods, []string{"notifications/initialized"}) {
+				t.Errorf("summary = %+v, want two notifications", got)
+			}
+		default:
+			t.Error("OnRequestSummary was not called for decoded batch")
+		}
+	})
+
+	t.Run("legacy decoded batch accepted and observed once", func(t *testing.T) {
+		var summaries []StreamableHTTPRequestSummary
+		handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{
+			Stateless: true,
+			OnRequestSummary: func(_ context.Context, summary StreamableHTTPRequestSummary) {
+				summaries = append(summaries, summary)
+			},
+		})
+		req := newRequest(`[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"notifications/initialized"}]`)
+		req.Header.Set(protocolVersionHeader, protocolVersion20250326)
+		rec := httptest.NewRecorder()
+		handler.ServeHTTP(rec, req)
+		if rec.Code != http.StatusAccepted {
+			t.Errorf("status = %d, want %d (body=%q)", rec.Code, http.StatusAccepted, rec.Body.String())
+		}
+		if len(summaries) != 1 || summaries[0].BatchCount != 2 || !summaries[0].IsNotification {
+			t.Errorf("summaries = %+v, want one notification-only batch", summaries)
 		}
 	})