internal/mcp: add batching support

Update the jsonrpc2 MCP framer to support both incoming and outgoing
batches. In order to achieve this, we must correlate the framed Reader
and Writer, which is not explicitly supported by the jsonrpc2 API, but
does work. A note is left to revisit the Framer interface.

Change-Id: If060cb5822da067833db20d58f5f112a0528da91
Reviewed-on: https://go-review.googlesource.com/c/tools/+/667295
Reviewed-by: Jonathan Amsterdam <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: findleyr

internal/jsonrpc2_v2/frame.go

@@ -39,6 +39,15 @@ type Writer interface {
 // Framer wraps low level byte readers and writers into jsonrpc2 message
 // readers and writers.
 // It is responsible for the framing and encoding of messages into wire form.
+//
+// TODO(rfindley): rethink the framer interface, as with JSONRPC2 batching
+// there is a need for Reader and Writer to be correlated, and while the
+// implementation of framing here allows that, it is not made explicit by the
+// interface.
+//
+// Perhaps a better interface would be
+//
+//	Frame(io.ReadWriteCloser) (Reader, Writer).
 type Framer interface {
 	// Reader wraps a byte reader into a message reader.
 	Reader(io.Reader) Reader

internal/mcp/mcp.go

@@ -19,5 +19,4 @@
 //   - Support multiple versions of the spec.
 //   - Implement proper JSON schema support, with both client-side and
 //     server-side validation..
-//   - Support batched JSON messages.
 package mcp

internal/mcp/mcp_test.go

@@ -142,6 +142,7 @@ func TestServerClosing(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
@@ -155,7 +156,36 @@ func TestServerClosing(t *testing.T) {
 	}
 	cc.Close()
 	wg.Wait()
-	if _, err = sc.CallTool(ctx, "greet", hiParams{"user"}); !errors.Is(err, mcp.ErrConnectionClosed) {
+	if _, err := sc.CallTool(ctx, "greet", hiParams{"user"}); !errors.Is(err, mcp.ErrConnectionClosed) {
 		t.Errorf("after disconnection, got error %v, want EOF", err)
 	}
 }
+
+func TestBatching(t *testing.T) {
+	ctx := context.Background()
+	ct, st := mcp.NewLocalTransport()
+
+	s := mcp.NewServer("testServer", "v1.0.0", nil)
+	_, err := s.Connect(ctx, st, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	c := mcp.NewClient("testClient", "v1.0.0", nil)
+	opts := new(mcp.ConnectionOptions)
+	mcp.BatchSize(opts, 2)
+	sc, err := c.Connect(ctx, ct, opts)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer sc.Close()
+
+	errs := make(chan error, 2)
+	for range 2 {
+		go func() {
+			_, err := sc.ListTools(ctx)
+			errs <- err
+		}()
+	}
+
+}

internal/mcp/transport.go

@@ -13,6 +13,7 @@ import (
 	"log"
 	"net"
 	"os"
+	"sync"
 
 	jsonrpc2 "golang.org/x/tools/internal/jsonrpc2_v2"
 )
@@ -33,7 +34,9 @@ type Transport struct {
 // ConnectionOptions configures the behavior of an individual client<->server
 // connection.
 type ConnectionOptions struct {
-	Logger io.Writer
+	Logger io.Writer // if set, write RPC logs
+
+	batchSize int // outgoing batch size for requests/notifications, for testing
 }
 
 // NewStdIOTransport constructs a transport that communicates over
@@ -215,40 +218,238 @@ func (r rwc) Close() error {
 }
 
 // A ndjsonFramer is a jsonrpc2.Framer that delimits messages with newlines.
+// It also supports jsonrpc2 batching.
+//
+// See https://github.com/ndjson/ndjson-spec for discussion of newline
+// delimited JSON.
 //
-// See also https://github.com/ndjson/ndjson-spec.
-type ndjsonFramer struct{}
-type rawReader struct{ in *json.Decoder } // relies on json.Decoder message boundaries
-type ndjsonWriter struct{ out io.Writer } // writes newline message boundaries
+// See [msgBatch] for more discussion of message batching.
+type ndjsonFramer struct {
+	// batchSize allows customizing batching behavior for testing.
+	//
+	// If set to a positive number, requests and notifications will be buffered
+	// into groups of this size before being sent as a batch.
+	batchSize int
+
+	// batches correlate incoming requests to the batch in which they arrived.
+	batchMu sync.Mutex
+	batches map[jsonrpc2.ID]*msgBatch // lazily allocated
+}
 
-func (ndjsonFramer) Reader(rw io.Reader) jsonrpc2.Reader {
-	return &rawReader{in: json.NewDecoder(rw)}
+// addBatch records a msgBatch for an incoming batch payload.
+// It returns an error if batch is malformed, containing previously seen IDs.
+//
+// See [msgBatch] for more.
+func (f *ndjsonFramer) addBatch(batch *msgBatch) error {
+	f.batchMu.Lock()
+	defer f.batchMu.Unlock()
+	for id := range batch.unresolved {
+		if _, ok := f.batches[id]; ok {
+			return fmt.Errorf("%w: batch contains previously seen request %v", jsonrpc2.ErrInvalidRequest, id.Raw())
+		}
+	}
+	for id := range batch.unresolved {
+		if f.batches == nil {
+			f.batches = make(map[jsonrpc2.ID]*msgBatch)
+		}
+		f.batches[id] = batch
+	}
+	return nil
 }
 
-func (ndjsonFramer) Writer(rw io.Writer) jsonrpc2.Writer {
-	return &ndjsonWriter{out: rw}
+// updateBatch records a response in the message batch tracking the
+// corresponding incoming call, if any.
+//
+// The second result reports whether resp was part of a batch. If this is true,
+// the first result is nil if the batch is still incomplete, or the full set of
+// batch responses if resp completed the batch.
+func (f *ndjsonFramer) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, bool) {
+	f.batchMu.Lock()
+	defer f.batchMu.Unlock()
+
+	if batch, ok := f.batches[resp.ID]; ok {
+		idx, ok := batch.unresolved[resp.ID]
+		if !ok {
+			panic("internal error: inconsistent batches")
+		}
+		batch.responses[idx] = resp
+		delete(batch.unresolved, resp.ID)
+		delete(f.batches, resp.ID)
+		if len(batch.unresolved) == 0 {
+			return batch.responses, true
+		}
+		return nil, true
+	}
+	return nil, false
 }
 
-func (r *rawReader) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
+// A msgBatch records information about an incoming batch of JSONRPC2 calls.
+//
+// The JSONRPC2 spec (https://www.jsonrpc.org/specification#batch) says:
+//
+// "The Server should respond with an Array containing the corresponding
+// Response objects, after all of the batch Request objects have been
+// processed. A Response object SHOULD exist for each Request object, except
+// that there SHOULD NOT be any Response objects for notifications. The Server
+// MAY process a batch rpc call as a set of concurrent tasks, processing them
+// in any order and with any width of parallelism."
+//
+// Therefore, a msgBatch keeps track of outstanding calls and their responses.
+// When there are no unresolved calls, the response payload is sent.
+type msgBatch struct {
+	unresolved map[jsonrpc2.ID]int
+	responses  []*jsonrpc2.Response
+}
+
+// An ndjsonReader reads newline-delimited messages or message batches.
+type ndjsonReader struct {
+	queue  []jsonrpc2.Message
+	framer *ndjsonFramer
+	in     *json.Decoder
+}
+
+// A ndjsonWriter writes newline-delimited messages to the wrapped io.Writer.
+//
+// If batch is set, messages are wrapped in a JSONRPC2 batch.
+type ndjsonWriter struct {
+	// Testing support: if outgoingBatch has capacity, it is used to buffer
+	// outgoing messages before sending a JSONRPC2 message batch.
+	outgoingBatch []jsonrpc2.Message
+
+	framer *ndjsonFramer // to track batch responses
+	out    io.Writer     // to write to the wire
+}
+
+func (f *ndjsonFramer) Reader(r io.Reader) jsonrpc2.Reader {
+	return &ndjsonReader{framer: f, in: json.NewDecoder(r)}
+}
+
+func (f *ndjsonFramer) Writer(w io.Writer) jsonrpc2.Writer {
+	writer := &ndjsonWriter{framer: f, out: w}
+	if f.batchSize > 0 {
+		writer.outgoingBatch = make([]jsonrpc2.Message, 0, f.batchSize)
+	}
+	return writer
+}
+
+func (r *ndjsonReader) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
 	select {
 	case <-ctx.Done():
 		return nil, 0, ctx.Err()
 	default:
 	}
+	if len(r.queue) > 0 {
+		next := r.queue[0]
+		r.queue = r.queue[1:]
+		return next, 0, nil
+	}
 	var raw json.RawMessage
 	if err := r.in.Decode(&raw); err != nil {
 		return nil, 0, err
 	}
+	var rawBatch []json.RawMessage
+	if err := json.Unmarshal(raw, &rawBatch); err == nil {
+		msg, err := r.readBatch(rawBatch)
+		if err != nil {
+			return nil, 0, err
+		}
+		return msg, int64(len(raw)), nil
+	}
 	msg, err := jsonrpc2.DecodeMessage(raw)
 	return msg, int64(len(raw)), err
 }
 
+// readBatch reads a batch of jsonrpc2 messages, and records the batch
+// in the framer so that responses can be collected and send back together.
+func (r *ndjsonReader) readBatch(rawBatch []json.RawMessage) (jsonrpc2.Message, error) {
+	if len(rawBatch) == 0 {
+		return nil, fmt.Errorf("empty batch")
+	}
+
+	// From the spec:
+	// "If the batch rpc call itself fails to be recognized as an valid JSON or
+	// as an Array with at least one value, the response from the Server MUST be
+	// a single Response object. If there are no Response objects contained
+	// within the Response array as it is to be sent to the client, the server
+	// MUST NOT return an empty Array and should return nothing at all."
+	//
+	// In our case, an error actually breaks the jsonrpc2 connection entirely,
+	// but defensively we collect batch information before recording it, so that
+	// we don't leave the framer in an inconsistent state.
+	var (
+		first     jsonrpc2.Message   // first message, to return
+		queue     []jsonrpc2.Message // remaining messages
+		respBatch *msgBatch          // tracks incoming requests in the batch
+	)
+	for i, raw := range rawBatch {
+		msg, err := jsonrpc2.DecodeMessage(raw)
+		if err != nil {
+			return nil, err
+		}
+		if i == 0 {
+			first = msg
+		} else {
+			queue = append(queue, msg)
+		}
+		if req, ok := msg.(*jsonrpc2.Request); ok {
+			if respBatch == nil {
+				respBatch = &msgBatch{
+					unresolved: make(map[jsonrpc2.ID]int),
+				}
+			}
+			respBatch.unresolved[req.ID] = len(respBatch.responses)
+			respBatch.responses = append(respBatch.responses, nil)
+		}
+	}
+	if respBatch != nil {
+		// The batch contains one or more incoming requests to track.
+		if err := r.framer.addBatch(respBatch); err != nil {
+			return nil, err
+		}
+	}
+
+	r.queue = append(r.queue, queue...)
+	return first, nil
+}
+
 func (w *ndjsonWriter) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
 	select {
 	case <-ctx.Done():
 		return 0, ctx.Err()
 	default:
 	}
+
+	// Batching support: if msg is a Response, it may have completed a batch, so
+	// check that first. Otherwise, it is a request or notification, and we may
+	// want to collect it into a batch before sending, if we're configured to use
+	// outgoing batches.
+	if resp, ok := msg.(*jsonrpc2.Response); ok {
+		if batch, ok := w.framer.updateBatch(resp); ok {
+			if len(batch) > 0 {
+				data, err := marshalMessages(batch)
+				if err != nil {
+					return 0, err
+				}
+				data = append(data, '\n')
+				n, err := w.out.Write(data)
+				return int64(n), err
+			}
+			return 0, nil
+		}
+	} else if len(w.outgoingBatch) < cap(w.outgoingBatch) {
+		w.outgoingBatch = append(w.outgoingBatch, msg)
+		if len(w.outgoingBatch) == cap(w.outgoingBatch) {
+			data, err := marshalMessages(w.outgoingBatch)
+			w.outgoingBatch = w.outgoingBatch[:0]
+			if err != nil {
+				return 0, err
+			}
+			data = append(data, '\n')
+			n, err := w.out.Write(data)
+			return int64(n), err
+		}
+		return 0, nil
+	}
 	data, err := jsonrpc2.EncodeMessage(msg)
 	if err != nil {
 		return 0, fmt.Errorf("marshaling message: %v", err)
@@ -257,3 +458,15 @@ func (w *ndjsonWriter) Write(ctx context.Context, msg jsonrpc2.Message) (int64,
 	n, err := w.out.Write(data)
 	return int64(n), err
 }
+
+func marshalMessages[T jsonrpc2.Message](msgs []T) ([]byte, error) {
+	var rawMsgs []json.RawMessage
+	for _, msg := range msgs {
+		raw, err := jsonrpc2.EncodeMessage(msg)
+		if err != nil {
+			return nil, fmt.Errorf("encoding batch message: %w", err)
+		}
+		rawMsgs = append(rawMsgs, raw)
+	}
+	return json.Marshal(rawMsgs)
+}