From 80565a9dfd249f9711b2a0c0460f258ebae015b6 Mon Sep 17 00:00:00 2001 From: tangerg Date: Thu, 5 Nov 2026 17:32:10 +0800 Subject: [PATCH 1/2] mcp: let a canceled ioConn write leave the queue ioConn.Write serialized writers with a mutex and checked ctx only before waiting for it, so a write whose context was canceled while another write was in progress still waited for that write and was then sent anyway. A write blocked on a peer that stopped reading could not be interrupted at all. Queue writers on a one-slot channel so a canceled waiter returns without sending. Canceling the active write closes the connection: bytes already handed to the stream cannot be withdrawn, and a partial frame would corrupt every later message. Co-Authored-By: Claude Opus 5.5 --- mcp/transport.go | 41 +++++++++++------- mcp/transport_test.go | 96 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 14 deletions(-) diff --git a/mcp/transport.go b/mcp/transport.go index 724f134b..ab09ac67 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -547,7 +547,9 @@ type ioConn struct { protocolVersion string dropped map[jsonrpc2.ID]bool // lazily allocated - writeMu sync.Mutex // guards Write, which must be concurrency safe. + // sending is a one-slot send queue. Its holder alone writes rwc and + // outgoingBatch; waiting for it is cancelable, unlike a mutex. + sending chan struct{} rwc io.ReadWriteCloser // the underlying stream // incoming receives messages from the read loop started in [newIOConn]. @@ -639,6 +641,7 @@ func newIOConnLimited(rwc io.ReadWriteCloser, maxLineLength int) *ioConn { } }() return &ioConn{ + sending: make(chan struct{}, 1), rwc: rwc, incoming: incoming, closed: closed, @@ -869,15 +872,17 @@ func readBatch(data []byte) (msgs []jsonrpc.Message, isBatch bool, _ error) { } func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { - // As in [ioConn.Read], enforce that Writes on a closed context are an error. + // A message still waiting for the stream leaves the queue without being + // sent, and its cancellation does not disturb the sender ahead of it. select { + case t.sending <- struct{}{}: case <-ctx.Done(): return ctx.Err() - default: } - - t.writeMu.Lock() - defer t.writeMu.Unlock() + defer func() { <-t.sending }() + if err := ctx.Err(); err != nil { + return err + } // 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 @@ -893,9 +898,7 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { if err != nil { return err } - data = append(data, '\n') - _, err = t.rwc.Write(data) - return err + return t.send(ctx, append(data, '\n')) } return nil } @@ -907,9 +910,7 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { if err != nil { return err } - data = append(data, '\n') - _, err = t.rwc.Write(data) - return err + return t.send(ctx, append(data, '\n')) } return nil } @@ -917,8 +918,20 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { if err != nil { return fmt.Errorf("marshaling message: %v", err) } - data = append(data, '\n') // newline delimited - _, err = t.rwc.Write(data) + return t.send(ctx, append(data, '\n')) // newline delimited +} + +// send writes one complete frame. Bytes already handed to the stream cannot be +// withdrawn, and a partially written frame would corrupt every later one, so +// canceling ctx mid-write closes the connection. That interrupts the write only +// if closing rwc unblocks it; the caller then still learns that the peer may +// have received some or all of the frame. +func (t *ioConn) send(ctx context.Context, data []byte) error { + interrupt := context.AfterFunc(ctx, func() { _ = t.Close() }) + _, err := t.rwc.Write(data) + if !interrupt() { + return errors.Join(ctx.Err(), err) + } return err } diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 5852e3e3..e18bff74 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -12,7 +12,10 @@ import ( "fmt" "io" "log/slog" + "net" + "os" "strings" + "sync" "testing" "time" @@ -886,3 +889,96 @@ func TestLoggingConnForwardsOptionalInterfaces(t *testing.T) { t.Error("propagateCancellation() = true for a delegate without the interface, want false") } } + +// writeObserver reports when the connection starts writing to its peer. +type writeObserver struct { + net.Conn + once sync.Once + writing chan struct{} +} + +func (w *writeObserver) Write(p []byte) (int, error) { + w.once.Do(func() { close(w.writing) }) + return w.Conn.Write(p) +} + +func observedPipe(t *testing.T) (*ioConn, *writeObserver, net.Conn) { + t.Helper() + local, peer := net.Pipe() + observed := &writeObserver{Conn: local, writing: make(chan struct{})} + conn := newIOConn(observed) + t.Cleanup(func() { + conn.Close() + peer.Close() + }) + return conn, observed, peer +} + +// TestIOConnCanceledQueuedWriteIsNeverSent checks that a message canceled while +// waiting behind a blocked sender returns promptly, never reaches the peer, and +// leaves the sender ahead of it intact. +func TestIOConnCanceledQueuedWriteIsNeverSent(t *testing.T) { + conn, observed, peer := observedPipe(t) + first := make(chan error, 1) + go func() { + first <- conn.Write(context.Background(), &jsonrpc.Request{ID: jsonrpc2.Int64ID(1), Method: "first"}) + }() + <-observed.writing + + ctx, cancel := context.WithCancel(context.Background()) + queued := make(chan error, 1) + go func() { + queued <- conn.Write(ctx, &jsonrpc.Request{ID: jsonrpc2.Int64ID(2), Method: "canceled"}) + }() + time.Sleep(20 * time.Millisecond) // let the second message reach the queue + cancel() + select { + case err := <-queued: + if !errors.Is(err, context.Canceled) { + t.Fatalf("queued Write() = %v, want context.Canceled", err) + } + case <-time.After(5 * time.Second): + t.Fatal("queued Write() did not return after cancellation") + } + + frames := bufio.NewReader(peer) + line, err := frames.ReadString('\n') + if err != nil || !strings.Contains(line, `"first"`) { + t.Fatalf("peer read %q, %v; want the first message", line, err) + } + if err := <-first; err != nil { + t.Fatalf("first Write() = %v, want nil", err) + } + peer.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + if line, err := frames.ReadString('\n'); !errors.Is(err, os.ErrDeadlineExceeded) { + t.Fatalf("peer read %q, %v after the canceled message; want nothing", line, err) + } +} + +// TestIOConnCanceledActiveWriteClosesConnection checks that canceling a write +// blocked by a peer that never reads returns the cancellation and closes the +// connection, since a partial frame may already be on the stream. +func TestIOConnCanceledActiveWriteClosesConnection(t *testing.T) { + conn, observed, _ := observedPipe(t) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- conn.Write(ctx, &jsonrpc.Request{ID: jsonrpc2.Int64ID(1), Method: "blocked"}) + }() + <-observed.writing + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("Write() = %v, want context.Canceled", err) + } + case <-time.After(5 * time.Second): + t.Fatal("Write() did not return after cancellation") + } + if _, err := conn.Read(context.Background()); err == nil { + t.Fatal("Read() succeeded on a connection whose write was interrupted") + } + if err := conn.Write(context.Background(), &jsonrpc.Request{ID: jsonrpc2.Int64ID(2), Method: "after"}); err == nil { + t.Fatal("Write() succeeded on a connection whose write was interrupted") + } +} From a78ab1f1340d8c3b6d21e3ceac16c13534826ee2 Mon Sep 17 00:00:00 2001 From: tangerg Date: Fri, 6 Nov 2026 09:57:02 +0800 Subject: [PATCH 2/2] mcp: make waiting for the ioConn write lock cancelable As discussed in review, this change no longer interrupts a write that is already in progress by closing the connection. It only stops a write from waiting for an earlier one once its context is done. writeMu keeps its name and becomes a 1-element semaphore channel, like reqHeaderMu in the net/http HTTP/2 client. A canceled waiter returns ctx.Err(), which jsonrpc2 already treats as a failure of the call rather than of the connection. The regression test uses testing/synctest instead of sleeping. Updates #1498 Co-Authored-By: Claude Opus 5.5 --- mcp/transport.go | 47 +++++++-------- mcp/transport_test.go | 132 ++++++++++++++---------------------------- 2 files changed, 65 insertions(+), 114 deletions(-) diff --git a/mcp/transport.go b/mcp/transport.go index ab09ac67..d475152b 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -547,9 +547,10 @@ type ioConn struct { protocolVersion string dropped map[jsonrpc2.ID]bool // lazily allocated - // sending is a one-slot send queue. Its holder alone writes rwc and - // outgoingBatch; waiting for it is cancelable, unlike a mutex. - sending chan struct{} + // writeMu is a 1-element semaphore guarding Write, which must be + // concurrency safe. Send to it to lock and receive from it to unlock. + // Unlike a sync.Mutex, waiting for it can stop when the context is done. + writeMu chan struct{} rwc io.ReadWriteCloser // the underlying stream // incoming receives messages from the read loop started in [newIOConn]. @@ -641,7 +642,7 @@ func newIOConnLimited(rwc io.ReadWriteCloser, maxLineLength int) *ioConn { } }() return &ioConn{ - sending: make(chan struct{}, 1), + writeMu: make(chan struct{}, 1), rwc: rwc, incoming: incoming, closed: closed, @@ -872,17 +873,19 @@ func readBatch(data []byte) (msgs []jsonrpc.Message, isBatch bool, _ error) { } func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { - // A message still waiting for the stream leaves the queue without being - // sent, and its cancellation does not disturb the sender ahead of it. + // As in [ioConn.Read], enforce that Writes on a closed context are an error. select { - case t.sending <- struct{}{}: case <-ctx.Done(): return ctx.Err() + default: } - defer func() { <-t.sending }() - if err := ctx.Err(); err != nil { - return err + + select { + case t.writeMu <- struct{}{}: + case <-ctx.Done(): + return ctx.Err() } + defer func() { <-t.writeMu }() // 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 @@ -898,7 +901,9 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { if err != nil { return err } - return t.send(ctx, append(data, '\n')) + data = append(data, '\n') + _, err = t.rwc.Write(data) + return err } return nil } @@ -910,7 +915,9 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { if err != nil { return err } - return t.send(ctx, append(data, '\n')) + data = append(data, '\n') + _, err = t.rwc.Write(data) + return err } return nil } @@ -918,20 +925,8 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { if err != nil { return fmt.Errorf("marshaling message: %v", err) } - return t.send(ctx, append(data, '\n')) // newline delimited -} - -// send writes one complete frame. Bytes already handed to the stream cannot be -// withdrawn, and a partially written frame would corrupt every later one, so -// canceling ctx mid-write closes the connection. That interrupts the write only -// if closing rwc unblocks it; the caller then still learns that the peer may -// have received some or all of the frame. -func (t *ioConn) send(ctx context.Context, data []byte) error { - interrupt := context.AfterFunc(ctx, func() { _ = t.Close() }) - _, err := t.rwc.Write(data) - if !interrupt() { - return errors.Join(ctx.Err(), err) - } + data = append(data, '\n') // newline delimited + _, err = t.rwc.Write(data) return err } diff --git a/mcp/transport_test.go b/mcp/transport_test.go index e18bff74..b6551cfa 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -13,10 +13,9 @@ import ( "io" "log/slog" "net" - "os" "strings" - "sync" "testing" + "testing/synctest" "time" "github.com/google/jsonschema-go/jsonschema" @@ -890,95 +889,52 @@ func TestLoggingConnForwardsOptionalInterfaces(t *testing.T) { } } -// writeObserver reports when the connection starts writing to its peer. -type writeObserver struct { - net.Conn - once sync.Once - writing chan struct{} -} - -func (w *writeObserver) Write(p []byte) (int, error) { - w.once.Do(func() { close(w.writing) }) - return w.Conn.Write(p) -} - -func observedPipe(t *testing.T) (*ioConn, *writeObserver, net.Conn) { - t.Helper() - local, peer := net.Pipe() - observed := &writeObserver{Conn: local, writing: make(chan struct{})} - conn := newIOConn(observed) - t.Cleanup(func() { - conn.Close() - peer.Close() - }) - return conn, observed, peer -} +// TestIOConnCanceledQueuedWrite checks that a Write canceled while it waits for +// another Write returns without sending its message, and that the connection +// stays usable. +func TestIOConnCanceledQueuedWrite(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + local, peer := net.Pipe() + conn := newIOConn(local) + defer conn.Close() + defer peer.Close() + + write := func(ctx context.Context, method string) <-chan error { + errc := make(chan error, 1) + go func() { errc <- conn.Write(ctx, &jsonrpc.Request{Method: method}) }() + return errc + } -// TestIOConnCanceledQueuedWriteIsNeverSent checks that a message canceled while -// waiting behind a blocked sender returns promptly, never reaches the peer, and -// leaves the sender ahead of it intact. -func TestIOConnCanceledQueuedWriteIsNeverSent(t *testing.T) { - conn, observed, peer := observedPipe(t) - first := make(chan error, 1) - go func() { - first <- conn.Write(context.Background(), &jsonrpc.Request{ID: jsonrpc2.Int64ID(1), Method: "first"}) - }() - <-observed.writing + // The peer is not reading yet, so the first Write blocks in the pipe. + first := write(context.Background(), "first") + synctest.Wait() - ctx, cancel := context.WithCancel(context.Background()) - queued := make(chan error, 1) - go func() { - queued <- conn.Write(ctx, &jsonrpc.Request{ID: jsonrpc2.Int64ID(2), Method: "canceled"}) - }() - time.Sleep(20 * time.Millisecond) // let the second message reach the queue - cancel() - select { - case err := <-queued: - if !errors.Is(err, context.Canceled) { - t.Fatalf("queued Write() = %v, want context.Canceled", err) + ctx, cancel := context.WithCancel(context.Background()) + canceled := write(ctx, "canceled") + synctest.Wait() + cancel() + if err := <-canceled; !errors.Is(err, context.Canceled) { + t.Fatalf("canceled Write() = %v, want %v", err, context.Canceled) } - case <-time.After(5 * time.Second): - t.Fatal("queued Write() did not return after cancellation") - } - frames := bufio.NewReader(peer) - line, err := frames.ReadString('\n') - if err != nil || !strings.Contains(line, `"first"`) { - t.Fatalf("peer read %q, %v; want the first message", line, err) - } - if err := <-first; err != nil { - t.Fatalf("first Write() = %v, want nil", err) - } - peer.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) - if line, err := frames.ReadString('\n'); !errors.Is(err, os.ErrDeadlineExceeded) { - t.Fatalf("peer read %q, %v after the canceled message; want nothing", line, err) - } -} - -// TestIOConnCanceledActiveWriteClosesConnection checks that canceling a write -// blocked by a peer that never reads returns the cancellation and closes the -// connection, since a partial frame may already be on the stream. -func TestIOConnCanceledActiveWriteClosesConnection(t *testing.T) { - conn, observed, _ := observedPipe(t) - ctx, cancel := context.WithCancel(context.Background()) - done := make(chan error, 1) - go func() { - done <- conn.Write(ctx, &jsonrpc.Request{ID: jsonrpc2.Int64ID(1), Method: "blocked"}) - }() - <-observed.writing - cancel() - select { - case err := <-done: - if !errors.Is(err, context.Canceled) { - t.Fatalf("Write() = %v, want context.Canceled", err) + second := write(context.Background(), "second") + frames := bufio.NewScanner(peer) + for _, want := range []string{"first", "second"} { + if !frames.Scan() { + t.Fatalf("peer read no frame, want %q: %v", want, frames.Err()) + } + msg, err := jsonrpc2.DecodeMessage(frames.Bytes()) + if err != nil { + t.Fatalf("decoding %q: %v", frames.Bytes(), err) + } + if got := msg.(*jsonrpc.Request).Method; got != want { + t.Fatalf("peer read %q, want %q", got, want) + } } - case <-time.After(5 * time.Second): - t.Fatal("Write() did not return after cancellation") - } - if _, err := conn.Read(context.Background()); err == nil { - t.Fatal("Read() succeeded on a connection whose write was interrupted") - } - if err := conn.Write(context.Background(), &jsonrpc.Request{ID: jsonrpc2.Int64ID(2), Method: "after"}); err == nil { - t.Fatal("Write() succeeded on a connection whose write was interrupted") - } + for _, errc := range []<-chan error{first, second} { + if err := <-errc; err != nil { + t.Errorf("Write() = %v, want nil", err) + } + } + }) }