[mcp] mcputils for streamable http

lib/utils/mcputils/http.go

@@ -0,0 +1,139 @@
+/*
+ * Teleport
+ * Copyright (C) 2025  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package mcputils
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"mime"
+	"net/http"
+
+	"github.com/gravitational/trace"
+	"github.com/mark3labs/mcp-go/mcp"
+
+	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/lib/utils"
+)
+
+// ServerMessageProcessor defines an interface that process JSON RPC responses
+// and notifications.
+type ServerMessageProcessor interface {
+	// ProcessResponse process a response and returns the message for client.
+	ProcessResponse(context.Context, *JSONRPCResponse) mcp.JSONRPCMessage
+	// ProcessNotification process a notification and returns the message for client.
+	ProcessNotification(context.Context, *JSONRPCNotification) mcp.JSONRPCMessage
+}
+
+// ReplaceHTTPResponse handles replacing the MCP server response for the
+// streamable HTTP transport.
+//
+// https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#streamable-http
+func ReplaceHTTPResponse(ctx context.Context, resp *http.Response, processer ServerMessageProcessor) error {
+	// Nothing to replace.
+	if resp.StatusCode != http.StatusOK || resp.ContentLength == 0 {
+		return nil
+	}
+
+	mediaType, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	switch mediaType {
+	case "application/json":
+		// Single response.
+		respBody, err := utils.ReadAtMost(resp.Body, teleport.MaxHTTPRequestSize)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		respFromServer, err := unmarshalResponse(string(respBody))
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		respToClient := processer.ProcessResponse(ctx, respFromServer)
+		respToClientAsBody, err := json.Marshal(respToClient)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		resp.Body = io.NopCloser(bytes.NewReader(respToClientAsBody))
+		return nil
+
+	case "text/event-stream":
+		// Multiple messages (response or notification) can be sent through SSE.
+		// Instead of reading all messages then replacing them, here we replace
+		// the body with a reader that process the event one a time.
+		resp.Body = &httpSSEResponseReplacer{
+			ctx:               ctx,
+			SSEResponseReader: NewSSEResponseReader(resp.Body),
+			processer:         processer,
+		}
+		return nil
+	default:
+		return trace.BadParameter("unsupported response type %s", mediaType)
+	}
+}
+
+type httpSSEResponseReplacer struct {
+	*SSEResponseReader
+	ctx       context.Context
+	processer ServerMessageProcessor
+	buf       []byte
+}
+
+func (r *httpSSEResponseReplacer) Read(p []byte) (int, error) {
+	if len(r.buf) != 0 {
+		n := copy(p, r.buf)
+		r.buf = r.buf[n:]
+		return n, nil
+	}
+
+	msg, err := r.ReadMessage(r.ctx)
+	if err != nil {
+		if utils.IsOKNetworkError(err) {
+			return 0, io.EOF
+		}
+		return 0, trace.Wrap(err)
+	}
+
+	var base BaseJSONRPCMessage
+	if err := json.Unmarshal([]byte(msg), &base); err != nil {
+		return 0, trace.Wrap(err)
+	}
+
+	var respToClient mcp.JSONRPCMessage
+	switch {
+	case base.IsResponse():
+		respToClient = r.processer.ProcessResponse(r.ctx, base.MakeResponse())
+	case base.IsNotification():
+		respToClient = r.processer.ProcessNotification(r.ctx, base.MakeNotification())
+	default:
+		return 0, trace.BadParameter("message is not a response or a notification")
+	}
+
+	respToSendAsBody, err := json.Marshal(respToClient)
+	if err != nil {
+		return 0, trace.Wrap(err)
+	}
+
+	// Convert to SSE.
+	r.buf = fmt.Appendf(r.buf, "event: message\ndata: %s\n\n", string(respToSendAsBody))
+	return r.Read(p)
+}

lib/utils/mcputils/http_test.go

@@ -0,0 +1,131 @@
+/*
+ * Teleport
+ * Copyright (C) 2025  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package mcputils
+
+import (
+	"context"
+	"maps"
+	"net/http"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/gravitational/trace"
+	mcpclient "github.com/mark3labs/mcp-go/client"
+	mcpclienttransport "github.com/mark3labs/mcp-go/client/transport"
+	"github.com/mark3labs/mcp-go/mcp"
+	mcpserver "github.com/mark3labs/mcp-go/server"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport/lib/utils/mcptest"
+)
+
+func TestReplaceHTTPResponse(t *testing.T) {
+	t.Parallel()
+	ctx := t.Context()
+
+	// Set up a server.
+	mcpServer := mcptest.NewServer()
+	httpServer := mcpserver.NewTestStreamableHTTPServer(mcpServer)
+	t.Cleanup(httpServer.Close)
+
+	// Set up a client with custom transport which calls "ReplaceHTTPResponse".
+	httpClientTransport := newTestReplaceHTTPResponseTransport()
+	mcpClientTransport, err := mcpclienttransport.NewStreamableHTTP(
+		httpServer.URL+"/mcp",
+		mcpclienttransport.WithHTTPBasicClient(
+			&http.Client{Transport: httpClientTransport},
+		),
+		mcpclienttransport.WithContinuousListening(),
+	)
+	require.NoError(t, err)
+	client := mcpclient.NewClient(mcpClientTransport)
+	require.NoError(t, client.Start(ctx))
+
+	// Initialize client and call a tool.
+	_, err = mcptest.InitializeClient(ctx, client)
+	require.NoError(t, err)
+	mcptest.MustCallServerTool(t, ctx, client)
+	assert.Equal(t, uint32(2), httpClientTransport.countMCPResponse.Load())
+
+	// Send notifications from server. Notifications will be sent through SSE.
+	require.EventuallyWithT(t, func(collect *assert.CollectT) {
+		assert.Greater(collect, httpClientTransport.getCountMethods()["GET"], 0)
+	}, 2*time.Second, 100*time.Millisecond, "client SSE connected")
+	mcpServer.SendNotificationToAllClients("notifications/test", nil)
+	mcpServer.SendNotificationToAllClients("notifications/test", nil)
+	require.EventuallyWithT(t, func(collect *assert.CollectT) {
+		assert.Equal(collect, uint32(2), httpClientTransport.countMCPNotification.Load())
+	}, 2*time.Second, 100*time.Millisecond, "expected to receive notification")
+
+	// Close client and count the requests.
+	require.NoError(t, client.Close())
+	require.Equal(t, map[string]int{
+		"GET":    1, // For listening on SSE events.
+		"POST":   3, // "initialize", "notifications/initialize", and "tools/call".
+		"DELETE": 1, // Close session.
+	}, httpClientTransport.getCountMethods())
+}
+
+type testReplaceHTTPResponseTransport struct {
+	countMethods         map[string]int
+	countMethodsMu       sync.Mutex
+	countMCPResponse     atomic.Uint32
+	countMCPNotification atomic.Uint32
+}
+
+func newTestReplaceHTTPResponseTransport() *testReplaceHTTPResponseTransport {
+	return &testReplaceHTTPResponseTransport{
+		countMethods: make(map[string]int),
+	}
+}
+
+func (t *testReplaceHTTPResponseTransport) getCountMethods() map[string]int {
+	t.countMethodsMu.Lock()
+	defer t.countMethodsMu.Unlock()
+	return maps.Clone(t.countMethods)
+}
+
+func (t *testReplaceHTTPResponseTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	t.countMethodsMu.Lock()
+	t.countMethods[r.Method]++
+	t.countMethodsMu.Unlock()
+
+	resp, err := http.DefaultClient.Do(r)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if err := ReplaceHTTPResponse(r.Context(), resp, t); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	return resp, nil
+}
+
+func (t *testReplaceHTTPResponseTransport) ProcessResponse(_ context.Context, response *JSONRPCResponse) mcp.JSONRPCMessage {
+	t.countMCPResponse.Add(1)
+	return response
+}
+
+func (t *testReplaceHTTPResponseTransport) ProcessNotification(_ context.Context, notification *JSONRPCNotification) mcp.JSONRPCMessage {
+	t.countMCPNotification.Add(1)
+	return notification
+}

lib/utils/mcputils/protocol.go

@@ -56,13 +56,13 @@ func (p JSONRPCParams) GetName() (string, bool) {
 	return name, ok
 }
 
-// baseJSONRPCMessage is a base message that includes all fields for MCP
+// BaseJSONRPCMessage is a base message that includes all fields for MCP
 // protocol.
 //
 // Note that json.RawMessage is used to keep the original content when
 // marshaling it again. json.RawMessage can also be easily unmarshalled to user
 // defined types when needed. Same applies to other types in this file.
-type baseJSONRPCMessage struct {
+type BaseJSONRPCMessage struct {
 	// JSONRPC specifies the version of JSONRPC.
 	JSONRPC string `json:"jsonrpc"`
 	// ID is the ID for request and response. ID is nil for notification.
@@ -77,32 +77,42 @@ type baseJSONRPCMessage struct {
 	Error json.RawMessage `json:"error,omitempty"`
 }
 
-func (m *baseJSONRPCMessage) isNotification() bool {
+// IsNotification returns true if the message is a notification.
+func (m *BaseJSONRPCMessage) IsNotification() bool {
 	return m.ID.IsNil()
 }
-func (m *baseJSONRPCMessage) isRequest() bool {
+
+// IsRequest returns true if the message is a request.
+func (m *BaseJSONRPCMessage) IsRequest() bool {
 	return !m.ID.IsNil() && m.Method != ""
 }
-func (m *baseJSONRPCMessage) isResponse() bool {
+
+// IsResponse returns if the message is a response.
+func (m *BaseJSONRPCMessage) IsResponse() bool {
 	return !m.ID.IsNil() && (m.Result != nil || m.Error != nil)
 }
 
-func (m *baseJSONRPCMessage) makeNotification() *JSONRPCNotification {
+// MakeNotification converts the base message to JSONRPCNotification.
+func (m *BaseJSONRPCMessage) MakeNotification() *JSONRPCNotification {
 	return &JSONRPCNotification{
 		JSONRPC: m.JSONRPC,
 		Method:  m.Method,
 		Params:  m.Params,
 	}
 }
-func (m *baseJSONRPCMessage) makeRequest() *JSONRPCRequest {
+
+// MakeRequest converts the base message to JSONRPCRequest.
+func (m *BaseJSONRPCMessage) MakeRequest() *JSONRPCRequest {
 	return &JSONRPCRequest{
 		JSONRPC: m.JSONRPC,
 		ID:      m.ID,
 		Method:  m.Method,
 		Params:  m.Params,
 	}
 }
-func (m *baseJSONRPCMessage) makeResponse() *JSONRPCResponse {
+
+// MakeResponse converts the base message to JSONRPCResponse.
+func (m *BaseJSONRPCMessage) MakeResponse() *JSONRPCResponse {
 	return &JSONRPCResponse{
 		JSONRPC: m.JSONRPC,
 		ID:      m.ID,
@@ -164,6 +174,19 @@ func (r *JSONRPCResponse) GetInitializeResult() (*mcp.InitializeResult, error) {
 	return &result, nil
 }
 
+// unmarshalResponse is a helper that unmarshalls a raw message to an
+// JSONRPCResponse.
+func unmarshalResponse(rawMessage string) (*JSONRPCResponse, error) {
+	var base BaseJSONRPCMessage
+	if err := json.Unmarshal([]byte(rawMessage), &base); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	if !base.IsResponse() {
+		return nil, trace.BadParameter("message is not a response")
+	}
+	return base.MakeResponse(), nil
+}
+
 const (
 	// MethodNotificationInitialized defines the method used for "initialized"
 	// notification. This notification is sent by the client after it receives