From f3e872226b8fd9e76c9328366b8e0de850780bcc Mon Sep 17 00:00:00 2001 From: CaliLuke Date: Tue, 22 Sep 2026 12:14:49 -0700 Subject: [PATCH] mcp: summarize decoded streamable HTTP batches --- mcp/streamable.go | 64 ++++++++++++++++++++++------- mcp/streamable_test.go | 91 +++++++++++++++++++++++++++++++++++++++--- 2 files changed, 135 insertions(+), 20 deletions(-) diff --git a/mcp/streamable.go b/mcp/streamable.go index 9f8f2234..30c1d943 100644 --- a/mcp/streamable.go +++ b/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 // _, to be consistent with the typescript implementation. diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 7d4782bf..73d75fc7 100644 --- a/mcp/streamable_test.go +++ b/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) } })