diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index fbe0688..6f48c9b 100644 --- a/internal/jsonrpc2/conn.go +++ b/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 diff --git a/internal/jsonrpc2/frame.go b/internal/jsonrpc2/frame.go index d5bbe75..46fcc9d 100644 --- a/internal/jsonrpc2/frame.go +++ b/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) diff --git a/mcp/sse.go b/mcp/sse.go index cf44276..bdc4770 100644 --- a/mcp/sse.go +++ b/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 { diff --git a/mcp/transport.go b/mcp/transport.go index f2d5c72..ac778db 100644 --- a/mcp/transport.go +++ b/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