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
4 changed files with 102 additions and 167 deletions
Download patch file Download diff file Expand all files Collapse all files

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
}
}()