diff --git a/internal/jsonrpc2/frame.go b/internal/jsonrpc2/frame.go index b73fec1..788e729 100644 --- a/internal/jsonrpc2/frame.go +++ b/internal/jsonrpc2/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 diff --git a/mcp/mcp.go b/mcp/mcp.go index 61c988e..5de696b 100644 --- a/mcp/mcp.go +++ b/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 diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index f968cc4..56b8a7a 100644 --- a/mcp/mcp_test.go +++ b/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 + }() + } + +} diff --git a/mcp/transport.go b/mcp/transport.go index e82a325..e9ed8d3 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -13,6 +13,7 @@ import ( "log" "net" "os" + "sync" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" ) @@ -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 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 https://github.com/ndjson/ndjson-spec for discussion of newline +// delimited JSON. +// +// 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 -func (ndjsonFramer) Reader(rw io.Reader) jsonrpc2.Reader { - return &rawReader{in: json.NewDecoder(rw)} + // batches correlate incoming requests to the batch in which they arrived. + batchMu sync.Mutex + batches map[jsonrpc2.ID]*msgBatch // lazily allocated } -func (ndjsonFramer) Writer(rw io.Writer) jsonrpc2.Writer { - return &ndjsonWriter{out: 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 (r *rawReader) Read(ctx context.Context) (jsonrpc2.Message, int64, error) { +// 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 +} + +// 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) +} diff --git a/mcp/transport_test.go b/mcp/transport_test.go new file mode 100644 index 0000000..f7f5ca0 --- /dev/null +++ b/mcp/transport_test.go @@ -0,0 +1,61 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "io" + "testing" + + "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" +) + +// BatchSize causes a connection to collect n requests or notifications before +// sending a batch on the wire (responses are always sent in isolation). +// +// Exported for testing in the mcp_test package. +func BatchSize(opts *ConnectionOptions, n int) { + opts.batchSize = n +} + +func TestBatchFraming(t *testing.T) { + // This test checks that the ndjsonFramer can read and write JSON batches. + // + // The framer is configured to write a batch size of 2, and we confirm that + // nothing is sent over the wire until the second write, at which point both + // messages become available. + ctx := context.Background() + + r, w := io.Pipe() + framer := ndjsonFramer{batchSize: 2} + reader := framer.Reader(r) + writer := framer.Writer(w) + + // Read the two messages into a channel, for easy testing later. + read := make(chan jsonrpc2.Message) + go func() { + for range 2 { + msg, _, _ := reader.Read(ctx) + read <- msg + } + }() + + // The first write should not yet be observed by the reader. + writer.Write(ctx, &jsonrpc2.Request{ID: jsonrpc2.Int64ID(1), Method: "test"}) + select { + case got := <-read: + t.Fatalf("after one write, got message %v", got) + default: + } + + // ...but the second write causes both messages to be observed. + writer.Write(ctx, &jsonrpc2.Request{ID: jsonrpc2.Int64ID(2), Method: "test"}) + for _, want := range []int64{1, 2} { + got := <-read + if got := got.(*jsonrpc2.Request).ID.Raw(); got != want { + t.Errorf("got message #%d, want #%d", got, want) + } + } +}