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.

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