modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/jsonrpc2_v2: refactor Reader/Writer and remove event dependency

Browse source 
In preparation for extracting the jsonrpc2_v2 package, update the Reader
and Writer interfaces to not report byte counts, and remove the
internal/event dependency.

Change-Id: Ie6b373736fe297c8427cd4ece7ddad33e275864b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/678956
Reviewed-by: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Rob Findley 2025-06-04 23:51:25 +00:00 • committed by Robert Findley
commit 82be9ce756
6 changed files with 102 additions and 167 deletions
Download patch file Download diff file Expand all files Collapse all files

87
internal/jsonrpc2/conn.go
View file

@ -13,11 +13,6 @@ import (
"sync"
"sync/atomic"
"time"
"golang.org/x/tools/internal/event"
"golang.org/x/tools/internal/event/keys"
"golang.org/x/tools/internal/event/label"
"golang.org/x/tools/internal/jsonrpc2"
)
// Binder builds a connection configuration.
@ -193,7 +188,6 @@ type incomingRequest struct {
*Request // the request being processed
ctx context.Context
cancel context.CancelFunc
endSpan func() // called (and set to nil) when the response is sent
}
// Bind returns the options unmodified.
@ -293,15 +287,9 @@ func (c *Connection) start(ctx context.Context, reader Reader, preempter Preempt
// The params will be marshaled to JSON before sending over the wire, and will
// be handed to the method invoked.
func (c *Connection) Notify(ctx context.Context, method string, params any) (err error) {
ctx, done := event.Start(ctx, method,
jsonrpc2.Method.Of(method),
jsonrpc2.RPCDirection.Of(jsonrpc2.Outbound),
)
attempted := false
defer func() {
labelStatus(ctx, err)
done()
if attempted {
c.updateInFlight(func(s *inFlightState) {
s.outgoingNotifications--
@ -332,7 +320,6 @@ func (c *Connection) Notify(ctx context.Context, method string, params any) (err
return fmt.Errorf("marshaling notify parameters: %v", err)
}
event.Metric(ctx, jsonrpc2.Started.Of(1))
return c.write(ctx, notify)
}
@ -344,17 +331,10 @@ func (c *Connection) Notify(ctx context.Context, method string, params any) (err
func (c *Connection) Call(ctx context.Context, method string, params any) *AsyncCall {
// Generate a new request identifier.
id := Int64ID(atomic.AddInt64(&c.seq, 1))
ctx, endSpan := event.Start(ctx, method,
jsonrpc2.Method.Of(method),
jsonrpc2.RPCDirection.Of(jsonrpc2.Outbound),
jsonrpc2.RPCID.Of(fmt.Sprintf("%q", id)),
)
ac := &AsyncCall{
id: id,
ready: make(chan struct{}),
ctx: ctx,
endSpan: endSpan,
id: id,
ready: make(chan struct{}),
}
// When this method returns, either ac is retired, or the request has been
// written successfully and the call is awaiting a response (to be provided by
@ -381,7 +361,6 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async
return ac
}
event.Metric(ctx, jsonrpc2.Started.Of(1))
if err := c.write(ctx, call); err != nil {
// Sending failed. We will never get a response, so deliver a fake one if it
// wasn't already retired by the connection breaking.
@ -400,10 +379,8 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async
type AsyncCall struct {
id ID
ready chan struct{} // closed after response has been set and span has been ended
ready chan struct{} // closed after response has been set
response *Response
ctx context.Context // for event logging only
endSpan func() // close the tracing span when all processing for the message is complete
}
// ID used for this call.
@ -431,12 +408,6 @@ func (ac *AsyncCall) retire(response *Response) {
}
ac.response = response
labelStatus(ac.ctx, response.Error)
ac.endSpan()
// Allow the trace context, which may retain a lot of reachable values,
// to be garbage-collected.
ac.ctx, ac.endSpan = nil, nil
close(ac.ready)
}
@ -525,18 +496,15 @@ func (c *Connection) Close() error {
func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter Preempter) {
var err error
for {
var (
msg Message
n int64
)
msg, n, err = reader.Read(ctx)
var msg Message
msg, err = reader.Read(ctx)
if err != nil {
break
}
switch msg := msg.(type) {
case *Request:
c.acceptRequest(ctx, msg, n, preempter)
c.acceptRequest(ctx, msg, preempter)
case *Response:
c.updateInFlight(func(s *inFlightState) {
@ -568,28 +536,14 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter
// acceptRequest either handles msg synchronously or enqueues it to be handled
// asynchronously.
func (c *Connection) acceptRequest(ctx context.Context, msg *Request, msgBytes int64, preempter Preempter) {
// Add a span to the context for this request.
labels := append(make([]label.Label, 0, 3), // Make space for the ID if present.
jsonrpc2.Method.Of(msg.Method),
jsonrpc2.RPCDirection.Of(jsonrpc2.Inbound),
)
if msg.IsCall() {
labels = append(labels, jsonrpc2.RPCID.Of(fmt.Sprintf("%q", msg.ID)))
}
ctx, endSpan := event.Start(ctx, msg.Method, labels...)
event.Metric(ctx,
jsonrpc2.Started.Of(1),
jsonrpc2.ReceivedBytes.Of(msgBytes))
func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter Preempter) {
// In theory notifications cannot be cancelled, but we build them a cancel
// context anyway.
ctx, cancel := context.WithCancel(ctx)
reqCtx, cancel := context.WithCancel(ctx)
req := &incomingRequest{
Request: msg,
ctx: ctx,
ctx: reqCtx,
cancel: cancel,
endSpan: endSpan,
}
// If the request is a call, add it to the incoming map so it can be
@ -722,10 +676,6 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
err = fmt.Errorf("%w: %q", ErrMethodNotFound, req.Method)
}
if req.endSpan == nil {
return c.internalErrorf("%#v produced a duplicate %q Response", from, req.Method)
}
if result != nil && err != nil {
c.internalErrorf("%#v returned a non-nil result with a non-nil error for %s:\n%v\n%#v", from, req.Method, err, result)
result = nil // Discard the spurious result and respond with err.
@ -761,16 +711,11 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
if err != nil {
// TODO: can/should we do anything with this error beyond writing it to the event log?
// (Is this the right label to attach to the log?)
event.Label(req.ctx, keys.Err.Of(err))
}
}
labelStatus(req.ctx, err)
// Cancel the request and finalize the event span to free any associated resources.
// Cancel the request to free any associated resources.
req.cancel()
req.endSpan()
req.endSpan = nil
c.updateInFlight(func(s *inFlightState) {
if s.incoming == 0 {
panic("jsonrpc2_v2: processResult called when incoming count is already zero")
@ -785,8 +730,7 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
func (c *Connection) write(ctx context.Context, msg Message) error {
writer := <-c.writer
defer func() { c.writer <- writer }()
n, err := writer.Write(ctx, msg)
event.Metric(ctx, jsonrpc2.SentBytes.Of(n))
err := writer.Write(ctx, msg)
if err != nil && ctx.Err() == nil {
// The call to Write failed, and since ctx.Err() is nil we can't attribute
@ -823,15 +767,6 @@ func (c *Connection) internalErrorf(format string, args ...any) error {
return fmt.Errorf("%w: %v", ErrInternal, err)
}
// labelStatus labels the status of the event in ctx based on whether err is nil.
func labelStatus(ctx context.Context, err error) {
if err == nil {
event.Label(ctx, jsonrpc2.StatusCode.Of("OK"))
} else {
event.Label(ctx, jsonrpc2.StatusCode.Of("ERROR"))
}
}
// notDone is a context.Context wrapper that returns a nil Done channel.
type notDone struct{ ctx context.Context }

76
internal/jsonrpc2/frame.go
View file

@ -22,7 +22,7 @@ import (
// a single Conn in a safe manner.
type Reader interface {
// Read gets the next message from the stream.
Read(context.Context) (Message, int64, error)
Read(context.Context) (Message, error)
}
// Writer abstracts the transport mechanics from the JSON RPC protocol.
@ -33,7 +33,7 @@ type Reader interface {
// a single Conn in a safe manner.
type Writer interface {
// Write sends a message to the stream.
Write(context.Context, Message) (int64, error)
Write(context.Context, Message) error
}
// Framer wraps low level byte readers and writers into jsonrpc2 message
@ -72,32 +72,32 @@ func (rawFramer) Writer(rw io.Writer) Writer {
return &rawWriter{out: rw}
}
func (r *rawReader) Read(ctx context.Context) (Message, int64, error) {
func (r *rawReader) Read(ctx context.Context) (Message, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
return nil, ctx.Err()
default:
}
var raw json.RawMessage
if err := r.in.Decode(&raw); err != nil {
return nil, 0, err
return nil, err
}
msg, err := DecodeMessage(raw)
return msg, int64(len(raw)), err
return msg, err
}
func (w *rawWriter) Write(ctx context.Context, msg Message) (int64, error) {
func (w *rawWriter) Write(ctx context.Context, msg Message) error {
select {
case <-ctx.Done():
return 0, ctx.Err()
return ctx.Err()
default:
}
data, err := EncodeMessage(msg)
if err != nil {
return 0, fmt.Errorf("marshaling message: %v", err)
return fmt.Errorf("marshaling message: %v", err)
}
n, err := w.out.Write(data)
return int64(n), err
_, err = w.out.Write(data)
return err
}
// HeaderFramer returns a new Framer.
@ -117,26 +117,29 @@ func (headerFramer) Writer(rw io.Writer) Writer {
return &headerWriter{out: rw}
}
func (r *headerReader) Read(ctx context.Context) (Message, int64, error) {
func (r *headerReader) Read(ctx context.Context) (Message, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
return nil, ctx.Err()
default:
}
var total, length int64
firstRead := true // to detect a clean EOF below
var contentLength int64
// read the header, stop on the first empty line
for {
line, err := r.in.ReadString('\n')
total += int64(len(line))
if err != nil {
if err == io.EOF {
if total == 0 {
return nil, 0, io.EOF
if firstRead && line == "" {
return nil, io.EOF // clean EOF
}
err = io.ErrUnexpectedEOF
}
return nil, total, fmt.Errorf("failed reading header line: %w", err)
return nil, fmt.Errorf("failed reading header line: %w", err)
}
firstRead = false
line = strings.TrimSpace(line)
// check we have a header line
if line == "" {
@ -144,49 +147,46 @@ func (r *headerReader) Read(ctx context.Context) (Message, int64, error) {
}
colon := strings.IndexRune(line, ':')
if colon < 0 {
return nil, total, fmt.Errorf("invalid header line %q", line)
return nil, fmt.Errorf("invalid header line %q", line)
}
name, value := line[:colon], strings.TrimSpace(line[colon+1:])
switch name {
case "Content-Length":
if length, err = strconv.ParseInt(value, 10, 32); err != nil {
return nil, total, fmt.Errorf("failed parsing Content-Length: %v", value)
if contentLength, err = strconv.ParseInt(value, 10, 32); err != nil {
return nil, fmt.Errorf("failed parsing Content-Length: %v", value)
}
if length <= 0 {
return nil, total, fmt.Errorf("invalid Content-Length: %v", length)
if contentLength <= 0 {
return nil, fmt.Errorf("invalid Content-Length: %v", contentLength)
}
default:
// ignoring unknown headers
}
}
if length == 0 {
return nil, total, fmt.Errorf("missing Content-Length header")
if contentLength == 0 {
return nil, fmt.Errorf("missing Content-Length header")
}
data := make([]byte, length)
n, err := io.ReadFull(r.in, data)
total += int64(n)
data := make([]byte, contentLength)
_, err := io.ReadFull(r.in, data)
if err != nil {
return nil, total, err
return nil, err
}
msg, err := DecodeMessage(data)
return msg, total, err
return msg, err
}
func (w *headerWriter) Write(ctx context.Context, msg Message) (int64, error) {
func (w *headerWriter) Write(ctx context.Context, msg Message) error {
select {
case <-ctx.Done():
return 0, ctx.Err()
return ctx.Err()
default:
}
data, err := EncodeMessage(msg)
if err != nil {
return 0, fmt.Errorf("marshaling message: %v", err)
return fmt.Errorf("marshaling message: %v", err)
}
n, err := fmt.Fprintf(w.out, "Content-Length: %v\r\n\r\n", len(data))
total := int64(n)
_, err = fmt.Fprintf(w.out, "Content-Length: %v\r\n\r\n", len(data))
if err == nil {
n, err = w.out.Write(data)
total += int64(n)
_, err = w.out.Write(data)
}
return total, err
return err
}

4
mcp/conformance_test.go
View file

@ -118,7 +118,7 @@ func runServerTest(t *testing.T, test *conformanceTest) {
}
writeMsg := func(msg jsonrpc2.Message) {
if _, err := cStream.Write(ctx, msg); err != nil {
if err := cStream.Write(ctx, msg); err != nil {
t.Fatalf("Write failed: %v", err)
}
}
@ -145,7 +145,7 @@ func runServerTest(t *testing.T, test *conformanceTest) {
// next incoming response.
nextResponse := func() (*jsonrpc2.Response, error, bool) {
for {
msg, _, err := cStream.Read(ctx)
msg, err := cStream.Read(ctx)
if err != nil {
// TODO(rfindley): we don't document (or want to document) that the in
// memory transports use a net.Pipe. How can users detect this failure?

46
mcp/sse.go
View file

@ -252,26 +252,26 @@ type sseServerStream struct {
}
// Read implements jsonrpc2.Reader.
func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
return nil, ctx.Err()
case msg := <-s.t.incoming:
return msg, 0, nil
return msg, nil
case <-s.t.done:
return nil, 0, io.EOF
return nil, io.EOF
}
}
// Write implements jsonrpc2.Writer.
func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
if ctx.Err() != nil {
return 0, ctx.Err()
return ctx.Err()
}
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil {
return 0, err
return err
}
s.t.mu.Lock()
@ -281,11 +281,11 @@ func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64
// exited, and so we must lock around this write and check isDone, which is
// set before the hanging GET exits.
if s.t.closed {
return 0, io.EOF
return io.EOF
}
n, err := writeEvent(s.t.w, event{name: "message", data: data})
return int64(n), err
_, err = writeEvent(s.t.w, event{name: "message", data: data})
return err
}
// Close implements io.Closer, and closes the session.
@ -451,48 +451,48 @@ func (c *sseClientStream) isDone() bool {
return c.closed
}
func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
return nil, ctx.Err()
case <-c.done:
return nil, 0, io.EOF
return nil, io.EOF
case data := <-c.incoming:
if c.isDone() {
return nil, 0, io.EOF
return nil, io.EOF
}
msg, err := jsonrpc2.DecodeMessage(data)
if err != nil {
return nil, 0, err
return nil, err
}
return msg, int64(len(data)), nil
return msg, nil
}
}
func (c *sseClientStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
func (c *sseClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil {
return 0, err
return err
}
if c.isDone() {
return 0, io.EOF
return io.EOF
}
req, err := http.NewRequestWithContext(ctx, "POST", c.msgEndpoint.String(), bytes.NewReader(data))
if err != nil {
return 0, err
return err
}
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return 0, err
return err
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return 0, fmt.Errorf("failed to write: %s", resp.Status)
return fmt.Errorf("failed to write: %s", resp.Status)
}
return int64(len(data)), nil
return nil
}
func (c *sseClientStream) Close() error {

54
mcp/transport.go
View file

@ -191,8 +191,8 @@ type loggingStream struct {
}
// loggingReader is a stream middleware that logs incoming messages.
func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
msg, n, err := s.delegate.Read(ctx)
func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
msg, err := s.delegate.Read(ctx)
if err != nil {
fmt.Fprintf(s.w, "read error: %v", err)
} else {
@ -202,12 +202,12 @@ func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, int64, erro
}
fmt.Fprintf(s.w, "read: %s\n", string(data))
}
return msg, n, err
return msg, err
}
// loggingWriter is a stream middleware that logs outgoing messages.
func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
n, err := s.delegate.Write(ctx, msg)
func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
err := s.delegate.Write(ctx, msg)
if err != nil {
fmt.Fprintf(s.w, "write error: %v", err)
} else {
@ -217,7 +217,7 @@ func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64,
}
fmt.Fprintf(s.w, "write: %s\n", string(data))
}
return n, err
return err
}
func (s *loggingStream) Close() error {
@ -345,35 +345,35 @@ type msgBatch struct {
responses []*jsonrpc2.Response
}
func (t *ioStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
func (t *ioStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
return t.read(ctx, t.in)
}
func (t *ioStream) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, int64, error) {
func (t *ioStream) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
return nil, ctx.Err()
default:
}
if len(t.queue) > 0 {
next := t.queue[0]
t.queue = t.queue[1:]
return next, 0, nil
return next, nil
}
var raw json.RawMessage
if err := in.Decode(&raw); err != nil {
return nil, 0, err
return nil, err
}
var rawBatch []json.RawMessage
if err := json.Unmarshal(raw, &rawBatch); err == nil {
msg, err := t.readBatch(rawBatch)
if err != nil {
return nil, 0, err
return nil, err
}
return msg, int64(len(raw)), nil
return msg, nil
}
msg, err := jsonrpc2.DecodeMessage(raw)
return msg, int64(len(raw)), err
return msg, err
}
// readBatch reads a batch of jsonrpc2 messages, and records the batch
@ -429,10 +429,10 @@ func (t *ioStream) readBatch(rawBatch []json.RawMessage) (jsonrpc2.Message, erro
return first, nil
}
func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
select {
case <-ctx.Done():
return 0, ctx.Err()
return ctx.Err()
default:
}
@ -445,13 +445,13 @@ func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, erro
if len(batch) > 0 {
data, err := marshalMessages(batch)
if err != nil {
return 0, err
return err
}
data = append(data, '\n')
n, err := t.rwc.Write(data)
return int64(n), err
_, err = t.rwc.Write(data)
return err
}
return 0, nil
return nil
}
} else if len(t.outgoingBatch) < cap(t.outgoingBatch) {
t.outgoingBatch = append(t.outgoingBatch, msg)
@ -459,21 +459,21 @@ func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, erro
data, err := marshalMessages(t.outgoingBatch)
t.outgoingBatch = t.outgoingBatch[:0]
if err != nil {
return 0, err
return err
}
data = append(data, '\n')
n, err := t.rwc.Write(data)
return int64(n), err
_, err = t.rwc.Write(data)
return err
}
return 0, nil
return nil
}
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil {
return 0, fmt.Errorf("marshaling message: %v", err)
return fmt.Errorf("marshaling message: %v", err)
}
data = append(data, '\n') // newline delimited
n, err := t.rwc.Write(data)
return int64(n), err
_, err = t.rwc.Write(data)
return err
}
func (t *ioStream) Close() error {

2
mcp/transport_test.go
View file

@ -38,7 +38,7 @@ func TestBatchFraming(t *testing.T) {
read := make(chan jsonrpc2.Message)
go func() {
for range 2 {
msg, _, _ := tport.Read(ctx)
msg, _ := tport.Read(ctx)
read <- msg
}
}()