internal/jsonrpc2: allow asynchronous writes

This CL introduces a fundamental change to the jsonrpc2 library:
connection writes, which were previously serialized by the jsonrpc2
library itself, are now allowed to be concurrent.

The change in semantics of the jsonrpc2 library should hopefully be easy
to review, since moving the synchronization to the Writer implementation
is equivalent to the previous logic. However, this change is critical
for the streamable client transport, because it allows for concurrent
http requests to the server.

Consider that a write is a POST to the server, and we don't know that
write succeeded until we get the response header. Previously, we had the
following problem: if the client POSTs a request, and the server blocks
its response on a request made through the hanging GET, the client was
unable to respond because the initial POST is still blocked. We could
update our streamable server transport to force a flush of the response
headers, but we can't guarantee that other servers behave the same way.
Fundamentally, writes in the spec are asynchronous, and we need to
support that.

internal/jsonrpc2/conn.go

@@ -65,8 +65,7 @@ type Connection struct {
 	state   inFlightState // accessed only in updateInFlight
 	done    chan struct{} // closed (under stateMu) when state.closed is true and all goroutines have completed
 
-	writer chan Writer // 1-buffered; stores the writer when not in use
-
+	writer  Writer
 	handler Handler
 
 	onInternalError func(error)
@@ -214,12 +213,11 @@ func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection {
 	c := &Connection{
 		state:           inFlightState{closer: cfg.Closer},
 		done:            make(chan struct{}),
-		writer:          make(chan Writer, 1),
+		writer:          cfg.Writer,
 		onDone:          cfg.OnDone,
 		onInternalError: cfg.OnInternalError,
 	}
 	c.handler = cfg.Bind(c)
-	c.writer <- cfg.Writer
 	c.start(ctx, cfg.Reader, cfg.Preempter)
 	return c
 }
@@ -239,7 +237,6 @@ func bindConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Bind
 	c := &Connection{
 		state:  inFlightState{closer: rwc},
 		done:   make(chan struct{}),
-		writer: make(chan Writer, 1),
 		onDone: onDone,
 	}
 	// It's tempting to set a finalizer on c to verify that the state has gone
@@ -259,7 +256,7 @@ func bindConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Bind
 	}
 	c.onInternalError = options.OnInternalError
 
-	c.writer <- framer.Writer(rwc)
+	c.writer = framer.Writer(rwc)
 	reader := framer.Reader(rwc)
 	c.start(ctx, reader, options.Preempter)
 	return c
@@ -728,9 +725,7 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
 // write is used by all things that write outgoing messages, including replies.
 // it makes sure that writes are atomic
 func (c *Connection) write(ctx context.Context, msg Message) error {
-	writer := <-c.writer
-	defer func() { c.writer <- writer }()
-	err := writer.Write(ctx, msg)
+	err := c.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

internal/jsonrpc2/frame.go

@@ -12,12 +12,14 @@ import (
 	"io"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 // Reader abstracts the transport mechanics from the JSON RPC protocol.
 // A Conn reads messages from the reader it was provided on construction,
 // and assumes that each call to Read fully transfers a single message,
 // or returns an error.
+//
 // A reader is not safe for concurrent use, it is expected it will be used by
 // a single Conn in a safe manner.
 type Reader interface {
@@ -29,8 +31,9 @@ type Reader interface {
 // A Conn writes messages using the writer it was provided on construction,
 // and assumes that each call to Write fully transfers a single message,
 // or returns an error.
-// A writer is not safe for concurrent use, it is expected it will be used by
-// a single Conn in a safe manner.
+//
+// A writer must be safe for concurrent use, as writes may occur concurrently
+// in practice: libraries may make calls or respond to requests asynchronously.
 type Writer interface {
 	// Write sends a message to the stream.
 	Write(context.Context, Message) error
@@ -62,7 +65,10 @@ func RawFramer() Framer { return rawFramer{} }
 
 type rawFramer struct{}
 type rawReader struct{ in *json.Decoder }
-type rawWriter struct{ out io.Writer }
+type rawWriter struct {
+	mu  sync.Mutex
+	out io.Writer
+}
 
 func (rawFramer) Reader(rw io.Reader) Reader {
 	return &rawReader{in: json.NewDecoder(rw)}
@@ -92,10 +98,14 @@ func (w *rawWriter) Write(ctx context.Context, msg Message) error {
 		return ctx.Err()
 	default:
 	}
+
 	data, err := EncodeMessage(msg)
 	if err != nil {
 		return fmt.Errorf("marshaling message: %v", err)
 	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
 	_, err = w.out.Write(data)
 	return err
 }
@@ -107,7 +117,10 @@ func HeaderFramer() Framer { return headerFramer{} }
 
 type headerFramer struct{}
 type headerReader struct{ in *bufio.Reader }
-type headerWriter struct{ out io.Writer }
+type headerWriter struct {
+	mu  sync.Mutex
+	out io.Writer
+}
 
 func (headerFramer) Reader(rw io.Reader) Reader {
 	return &headerReader{in: bufio.NewReader(rw)}
@@ -180,6 +193,9 @@ func (w *headerWriter) Write(ctx context.Context, msg Message) error {
 		return ctx.Err()
 	default:
 	}
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
 	data, err := EncodeMessage(msg)
 	if err != nil {
 		return fmt.Errorf("marshaling message: %v", err)

mcp/sse.go

@@ -155,7 +155,7 @@ func (t *SSEServerTransport) Connect(context.Context) (Connection, error) {
 	if err != nil {
 		return nil, err
 	}
-	return sseServerConn{t}, nil
+	return &sseServerConn{t: t}, nil
 }
 
 func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -244,10 +244,10 @@ type sseServerConn struct {
 }
 
 // TODO(jba): get the session ID. (Not urgent because SSE transports have been removed from the spec.)
-func (s sseServerConn) SessionID() string { return "" }
+func (s *sseServerConn) SessionID() string { return "" }
 
 // Read implements jsonrpc2.Reader.
-func (s sseServerConn) Read(ctx context.Context) (jsonrpc.Message, error) {
+func (s *sseServerConn) Read(ctx context.Context) (jsonrpc.Message, error) {
 	select {
 	case <-ctx.Done():
 		return nil, ctx.Err()
@@ -259,7 +259,7 @@ func (s sseServerConn) Read(ctx context.Context) (jsonrpc.Message, error) {
 }
 
 // Write implements jsonrpc2.Writer.
-func (s sseServerConn) Write(ctx context.Context, msg jsonrpc.Message) error {
+func (s *sseServerConn) Write(ctx context.Context, msg jsonrpc.Message) error {
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}
@@ -288,7 +288,7 @@ func (s sseServerConn) Write(ctx context.Context, msg jsonrpc.Message) error {
 // It must be safe to call Close more than once, as the close may
 // asynchronously be initiated by either the server closing its connection, or
 // by the hanging GET exiting.
-func (s sseServerConn) Close() error {
+func (s *sseServerConn) Close() error {
 	s.t.mu.Lock()
 	defer s.t.mu.Unlock()
 	if !s.t.closed {

mcp/transport.go

@@ -38,9 +38,25 @@ type Transport interface {
 
 // A Connection is a logical bidirectional JSON-RPC connection.
 type Connection interface {
+	// Read reads the next message to process off the connection.
+	//
+	// Read need not be safe for concurrent use: Read is called in a
+	// concurrency-safe manner by the JSON-RPC library.
 	Read(context.Context) (jsonrpc.Message, error)
+
+	// Write writes a new message to the connection.
+	//
+	// Write may be called concurrently, as calls or reponses may occur
+	// concurrently in user code.
 	Write(context.Context, jsonrpc.Message) error
-	Close() error // may be called concurrently by both peers
+
+	// Close closes the connection. It is implicitly called whenever a Read or
+	// Write fails.
+	//
+	// Close may be called multiple times, potentially concurrently.
+	Close() error
+
+	// TODO(#148): remove SessionID from this interface.
 	SessionID() string
 }
 
@@ -264,7 +280,8 @@ func (r rwc) Close() error {
 //
 // See [msgBatch] for more discussion of message batching.
 type ioConn struct {
-	rwc io.ReadWriteCloser // the underlying stream
+	writeMu sync.Mutex         // guards Write, which must be concurrency safe.
+	rwc     io.ReadWriteCloser // the underlying stream
 
 	// incoming receives messages from the read loop started in [newIOConn].
 	incoming <-chan msgOrErr
@@ -497,6 +514,9 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error {
 	default:
 	}
 
+	t.writeMu.Lock()
+	defer t.writeMu.Unlock()
+
 	// 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