diff --git a/mcp/streamable.go b/mcp/streamable.go index f642db1d..0a585895 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -69,6 +69,25 @@ type sessionInfo struct { timer *time.Timer } +// StreamableHTTPRequestSummary contains redacted metadata about a single +// JSON-RPC message 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. + Method string + + // RequestID is valid only for a 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. + IsNotification bool + + // IsResponse reports whether the message is a JSON-RPC response. + IsResponse bool +} + // startPOST signals that a POST request for this session is starting (which // carries a client->server message), pausing the session timeout if it was // running. @@ -218,6 +237,23 @@ type StreamableHTTPOptions struct { // Requests using older protocol versions (including those routed through // 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. + // 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 + // may contain values added by authentication or other middleware. + // + // The callback is not invoked for requests rejected before this point, + // including HTTP, authorization, session-routing, and connection failures, + // even if connection setup previously inspected the body. Use HTTP middleware + // to observe those failures. Use [Server.AddReceivingMiddleware] to observe + // dispatched messages; this callback additionally observes decoded messages + // that are rejected before dispatch without exposing their parameters. + OnRequestSummary func(context.Context, StreamableHTTPRequestSummary) } // DefaultMaxRequestBodyBytes is the default value used for @@ -422,14 +458,8 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http. } } - transport := &StreamableServerTransport{ - SessionID: sessionID, - Stateless: true, - EventStore: h.opts.EventStore, - jsonResponse: h.opts.JSONResponse, - logger: h.opts.Logger, - shouldPropagateCancellation: info.usesNewProtocol && (info.isSubscriptionsListen || h.opts.PropagateRequestCancellation), - } + transport := h.newStreamableServerTransport(sessionID, true) + transport.shouldPropagateCancellation = info.usesNewProtocol && (info.isSubscriptionsListen || h.opts.PropagateRequestCancellation) session, err := connectStreamable(req.Context(), server, transport, info.opts) if err != nil { @@ -534,6 +564,17 @@ func connectStreamable(ctx context.Context, server *Server, transport *Streamabl return s, nil } +func (h *StreamableHTTPHandler) newStreamableServerTransport(sessionID string, stateless bool) *StreamableServerTransport { + return &StreamableServerTransport{ + SessionID: sessionID, + Stateless: stateless, + EventStore: h.opts.EventStore, + jsonResponse: h.opts.JSONResponse, + logger: h.opts.Logger, + onRequestSummary: h.opts.OnRequestSummary, + } +} + // serveStateful handles requests for stateful servers. // Stateful servers support GET, POST, and DELETE, and maintain persistent // sessions keyed by session ID. @@ -652,13 +693,7 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht } sessionID = server.opts.GetSessionID() - transport := &StreamableServerTransport{ - SessionID: sessionID, - Stateless: false, - EventStore: h.opts.EventStore, - jsonResponse: h.opts.JSONResponse, - logger: h.opts.Logger, - } + transport := h.newStreamableServerTransport(sessionID, false) // Sessions without a session ID (GetSessionID returned "") are ephemeral: // there's no way to address them, so they are closed after the request. @@ -832,6 +867,10 @@ type StreamableServerTransport struct { // [streamableServerConn]. See its docstring. shouldPropagateCancellation bool + // onRequestSummary is forwarded from StreamableHTTPOptions for transports + // created by StreamableHTTPHandler. + onRequestSummary func(context.Context, StreamableHTTPRequestSummary) + // connection is non-nil if and only if the transport has been connected. connection *streamableServerConn } @@ -848,6 +887,7 @@ func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, er jsonResponse: t.jsonResponse, logger: ensureLogger(t.logger), // see #556: must be non-nil shouldPropagateCancellation: t.shouldPropagateCancellation, + onRequestSummary: t.onRequestSummary, incoming: make(chan jsonrpc.Message, 10), done: make(chan struct{}), streams: make(map[string]*stream), @@ -876,10 +916,11 @@ func (t *StreamableServerTransport) SupportsProtocolVersion(version string) bool } type streamableServerConn struct { - sessionID string - stateless bool - jsonResponse bool - eventStore EventStore + sessionID string + stateless bool + jsonResponse bool + eventStore EventStore + onRequestSummary func(context.Context, StreamableHTTPRequestSummary) // shouldPropagateCancellation is true when the underlying HTTP request's // lifetime IS the connection's cancellation signal (e.g., a stateless @@ -1437,6 +1478,9 @@ 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])) + } protocolVersion := protocolVersionFromContext(req.Context()) if protocolVersion == "" { @@ -1737,6 +1781,22 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques c.hangResponse(req.Context(), done) } +func summarizeStreamableHTTPRequest(msg jsonrpc.Message) StreamableHTTPRequestSummary { + var summary StreamableHTTPRequestSummary + switch msg := msg.(type) { + case *jsonrpc.Request: + summary.Method = msg.Method + if msg.IsCall() { + summary.RequestID = msg.ID + } else { + summary.IsNotification = true + } + case *jsonrpc.Response: + summary.IsResponse = true + } + 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 85ee331f..9c65e0dd 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -28,6 +28,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/iotest" "time" "github.com/google/go-cmp/cmp" @@ -3813,6 +3814,422 @@ func TestStreamableServerRejectsDuplicateInFlightRequestID(t *testing.T) { } } +func TestSummarizeStreamableHTTPRequest(t *testing.T) { + tests := []struct { + name string + body string + wantMethod string + wantRequestID any + wantIsNotification bool + wantIsResponse bool + }{ + { + name: "single integer ID call", + body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`, + wantMethod: "tools/list", + wantRequestID: int64(1), + }, + { + name: "single string ID call", + body: `{"jsonrpc":"2.0","id":"request-1","method":"ping"}`, + wantMethod: "ping", + wantRequestID: "request-1", + }, + { + name: "null ID is a notification", + body: `{"jsonrpc":"2.0","id":null,"method":"notifications/initialized"}`, + wantMethod: "notifications/initialized", + wantIsNotification: true, + }, + { + name: "response", + body: `{"jsonrpc":"2.0","id":3,"result":{}}`, + wantIsResponse: true, + }, + { + name: "empty notification method is retained", + body: `{"jsonrpc":"2.0","method":""}`, + wantIsNotification: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + msg, err := jsonrpc2.DecodeMessage([]byte(test.body)) + if err != nil { + t.Fatalf("DecodeMessage: %v", err) + } + got := summarizeStreamableHTTPRequest(msg) + if got.Method != test.wantMethod { + t.Errorf("Method = %q, want %q", got.Method, test.wantMethod) + } + if gotID := got.RequestID.Raw(); gotID != test.wantRequestID { + t.Errorf("RequestID.Raw() = %#v, want %#v", gotID, test.wantRequestID) + } + if got.IsNotification != test.wantIsNotification { + t.Errorf("IsNotification = %t, want %t", got.IsNotification, test.wantIsNotification) + } + if got.IsResponse != test.wantIsResponse { + t.Errorf("IsResponse = %t, want %t", got.IsResponse, test.wantIsResponse) + } + }) + } +} + +func TestStreamableHTTPRequestSummaryContext(t *testing.T) { + type contextKey struct{} + const contextValue = "middleware-value" + initBody := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}` + + for _, stateless := range []bool{false, true} { + t.Run(fmt.Sprintf("stateless=%t", stateless), func(t *testing.T) { + server := NewServer(testImpl, nil) + observed := make(chan StreamableHTTPRequestSummary, 1) + var summarySeen atomic.Bool + server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if !summarySeen.Load() { + t.Error("receiving middleware ran before OnRequestSummary") + } + if method != methodInitialize { + t.Errorf("receiving middleware method = %q, want %q", method, methodInitialize) + } + return next(ctx, method, req) + } + }) + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: stateless, + JSONResponse: true, + OnRequestSummary: func(ctx context.Context, summary StreamableHTTPRequestSummary) { + if got := ctx.Value(contextKey{}); got != contextValue { + t.Errorf("callback context value = %v, want %q", got, contextValue) + } + if summary.Method != methodInitialize { + t.Errorf("callback Method = %q, want %q", summary.Method, methodInitialize) + } + summarySeen.Store(true) + observed <- summary + }, + }) + defer handler.closeAll() + + req := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", strings.NewReader(initBody)) + req = req.WithContext(context.WithValue(req.Context(), contextKey{}, contextValue)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d (body=%q)", rec.Code, http.StatusOK, rec.Body.String()) + } + summary := <-observed + if summary.Method != methodInitialize { + t.Errorf("Method = %q, want %q", summary.Method, methodInitialize) + } + if got := summary.RequestID.Raw(); got != int64(1) { + t.Errorf("RequestID.Raw() = %#v, want 1", got) + } + if summary.IsNotification || summary.IsResponse { + t.Errorf("message kind = (notification=%t, response=%t), want call", summary.IsNotification, summary.IsResponse) + } + }) + } +} + +func TestStreamableHTTPRequestSummaryRejections(t *testing.T) { + newRequest := func(body string) *http.Request { + req := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + return req + } + + t.Run("decoded request rejected before dispatch is observed", func(t *testing.T) { + var calls atomic.Int64 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{ + Stateless: true, + OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) { + calls.Add(1) + }, + }) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, newRequest(`{"jsonrpc":"2.0","id":1,"method":"unknown/method"}`)) + if rec.Code != http.StatusBadRequest { + t.Errorf("status = %d, want %d", rec.Code, http.StatusBadRequest) + } + if got := calls.Load(); got != 1 { + t.Errorf("callback calls = %d, want 1", got) + } + }) + + t.Run("batch is not observed", func(t *testing.T) { + var calls atomic.Int64 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{ + Stateless: true, + OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) { + calls.Add(1) + }, + }) + req := newRequest(`[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"notifications/initialized"}]`) + req.Header.Set(protocolVersionHeader, protocolVersion20250618) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + 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) + } + }) + + t.Run("bodies not decoded are not observed", func(t *testing.T) { + tests := []struct { + name string + body io.ReadCloser + max int64 + }{ + {name: "empty", body: io.NopCloser(strings.NewReader(""))}, + {name: "malformed", body: io.NopCloser(strings.NewReader("{"))}, + {name: "unreadable", body: io.NopCloser(iotest.ErrReader(errors.New("read failed")))}, + {name: "over limit", body: io.NopCloser(strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`)), max: 8}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int64 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{ + Stateless: true, + MaxRequestBodyBytes: test.max, + OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) { + calls.Add(1) + }, + }) + req := newRequest("unused") + req.Body = test.body + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if got := calls.Load(); got != 0 { + t.Errorf("callback calls = %d, want 0 (status=%d)", got, rec.Code) + } + }) + } + }) + + t.Run("early session rejection is not observed", func(t *testing.T) { + var calls atomic.Int64 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{ + OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) { + calls.Add(1) + }, + }) + req := newRequest(`{"jsonrpc":"2.0","id":1,"method":"ping"}`) + req.Header.Set(sessionIDHeader, "missing-session") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusNotFound { + t.Errorf("status = %d, want %d", rec.Code, http.StatusNotFound) + } + if got := calls.Load(); got != 0 { + t.Errorf("callback calls = %d, want 0", got) + } + }) +} + +func TestStreamableHTTPRequestSummaryConcurrent(t *testing.T) { + const requestCount = 8 + allEntered := make(chan struct{}) + release := make(chan struct{}) + var entered atomic.Int64 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{ + Stateless: true, + OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) { + if entered.Add(1) == requestCount { + close(allEntered) + } + <-release + }, + }) + + statuses := make(chan int, requestCount) + for range requestCount { + go func() { + req := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"unknown/method"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + statuses <- rec.Code + }() + } + + concurrent := false + select { + case <-allEntered: + concurrent = true + case <-time.After(5 * time.Second): + } + close(release) + for range requestCount { + if status := <-statuses; status != http.StatusBadRequest { + t.Errorf("status = %d, want %d", status, http.StatusBadRequest) + } + } + if !concurrent { + t.Fatalf("only %d of %d callbacks ran concurrently", entered.Load(), requestCount) + } +} + +func TestStreamableHTTPRequestSummaryPanic(t *testing.T) { + panicValue := errors.New("observer panic") + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return NewServer(testImpl, nil) }, &StreamableHTTPOptions{ + Stateless: true, + OnRequestSummary: func(context.Context, StreamableHTTPRequestSummary) { + panic(panicValue) + }, + }) + req := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"unknown/method"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + + defer func() { + if got := recover(); got != panicValue { + t.Errorf("recovered panic = %v, want %v", got, panicValue) + } + }() + handler.ServeHTTP(httptest.NewRecorder(), req) +} + +type countingReader struct { + r io.Reader + reads int + bytes int +} + +func (r *countingReader) Read(p []byte) (int, error) { + r.reads++ + n, err := r.r.Read(p) + r.bytes += n + return n, err +} + +func TestStreamableHTTPRequestSummaryDoesNotChangeBodyHandling(t *testing.T) { + const body = `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}` + type result struct { + reads int + bytes int + bodyReplaced bool + status int + } + run := func(observe bool) result { + server := NewServer(testImpl, nil) + var routedReq *http.Request + opts := &StreamableHTTPOptions{Stateless: true, JSONResponse: true, MaxRequestBodyBytes: -1} + if observe { + opts.OnRequestSummary = func(context.Context, StreamableHTTPRequestSummary) {} + } + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { + routedReq = req + return server + }, opts) + reader := &countingReader{r: strings.NewReader(body)} + originalBody := io.NopCloser(reader) + req := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", nil) + req.Body = originalBody + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + return result{ + reads: reader.reads, + bytes: reader.bytes, + bodyReplaced: routedReq.Body != originalBody, + status: rec.Code, + } + } + + withoutObserver := run(false) + withObserver := run(true) + if diff := cmp.Diff(withoutObserver, withObserver, cmp.AllowUnexported(result{})); diff != "" { + t.Fatalf("enabling OnRequestSummary changed body handling (-without +with):\n%s", diff) + } + if withObserver.status != http.StatusOK { + t.Errorf("status = %d, want %d", withObserver.status, http.StatusOK) + } + if !withObserver.bodyReplaced { + t.Error("stateless pre-read baseline did not replace the routed request body") + } +} + +func TestStreamableHTTPRequestSummaryExistingSessionBodyHandling(t *testing.T) { + const initBody = `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}` + const pingBody = `{"jsonrpc":"2.0","id":2,"method":"ping"}` + type result struct { + reads int + bytes int + bodyReplaced bool + status int + } + run := func(observe bool) (result, int64) { + server := NewServer(testImpl, nil) + var observations atomic.Int64 + opts := &StreamableHTTPOptions{JSONResponse: true, MaxRequestBodyBytes: -1} + if observe { + opts.OnRequestSummary = func(context.Context, StreamableHTTPRequestSummary) { + observations.Add(1) + } + } + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts) + defer handler.closeAll() + + initReq := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", strings.NewReader(initBody)) + initReq.Header.Set("Content-Type", "application/json") + initReq.Header.Set("Accept", "application/json, text/event-stream") + initRec := httptest.NewRecorder() + handler.ServeHTTP(initRec, initReq) + if initRec.Code != http.StatusOK { + t.Fatalf("initialize status = %d, want %d (body=%q)", initRec.Code, http.StatusOK, initRec.Body.String()) + } + sessionID := initRec.Header().Get(sessionIDHeader) + if sessionID == "" { + t.Fatal("initialize response has no session ID") + } + + reader := &countingReader{r: strings.NewReader(pingBody)} + originalBody := io.NopCloser(reader) + req := httptest.NewRequest(http.MethodPost, "http://example.com/mcp", nil) + req.Body = originalBody + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(sessionIDHeader, sessionID) + req.Header.Set(protocolVersionHeader, protocolVersion20250618) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + return result{ + reads: reader.reads, + bytes: reader.bytes, + bodyReplaced: req.Body != originalBody, + status: rec.Code, + }, observations.Load() + } + + withoutObserver, withoutObservations := run(false) + withObserver, withObservations := run(true) + if diff := cmp.Diff(withoutObserver, withObserver, cmp.AllowUnexported(result{})); diff != "" { + t.Fatalf("enabling OnRequestSummary changed existing-session body handling (-without +with):\n%s", diff) + } + if withoutObservations != 0 { + t.Errorf("disabled observer calls = %d, want 0", withoutObservations) + } + if withObservations != 2 { + t.Errorf("enabled observer calls = %d, want 2", withObservations) + } + if withObserver.status != http.StatusOK { + t.Errorf("ping status = %d, want %d", withObserver.status, http.StatusOK) + } + if withObserver.bodyReplaced { + t.Error("existing-session request body was replaced") + } +} + // TestStreamableMaxRequestBodyBytes verifies that the streamable HTTP handler // enforces StreamableHTTPOptions.MaxRequestBodyBytes on incoming request // bodies. The limit must apply uniformly regardless of transfer encoding: