diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 7ad984b..3223945 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -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 } diff --git a/internal/jsonrpc2/frame.go b/internal/jsonrpc2/frame.go index 788e729..d5bbe75 100644 --- a/internal/jsonrpc2/frame.go +++ b/internal/jsonrpc2/frame.go @@ -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 } diff --git a/mcp/conformance_test.go b/mcp/conformance_test.go index abc88a0..6886d67 100644 --- a/mcp/conformance_test.go +++ b/mcp/conformance_test.go @@ -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? diff --git a/mcp/sse.go b/mcp/sse.go index 1c72e71..31dbb34 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -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 { diff --git a/mcp/transport.go b/mcp/transport.go index 1412627..91fcf10 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -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 { diff --git a/mcp/transport_test.go b/mcp/transport_test.go index b5be90f..51ec584 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -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 } }()