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
2 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
}