mcp: don't let a best-effort cancel notify delay connection teardown

internal/jsonrpc2/conn.go

@@ -51,6 +51,9 @@ 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
 
+	closing     chan struct{} // closed once when Close is first called
+	closingOnce sync.Once
+
 	writer  Writer
 	handler Handler
 
@@ -236,6 +239,7 @@ func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection {
 	c := &Connection{
 		state:           inFlightState{closer: cfg.Closer},
 		done:            make(chan struct{}),
+		closing:         make(chan struct{}),
 		writer:          cfg.Writer,
 		onDone:          cfg.OnDone,
 		onInternalError: cfg.OnInternalError,
@@ -502,12 +506,23 @@ func (c *Connection) wait(fromWait bool) error {
 // with IDs will receive immediate responses with ErrServerClosing, and no new
 // requests (not even notifications!) will be enqueued to the Handler.
 func (c *Connection) Close() error {
+	// Signal that teardown has begun so background best-effort work (e.g. a
+	// cancellation notify to an unresponsive peer) can stop instead of blocking
+	// the wait below.
+	c.closingOnce.Do(func() { close(c.closing) })
 	// Stop handling new requests, and interrupt the reader (by closing the
 	// connection) as soon as the active requests finish.
 	c.updateInFlight(func(s *inFlightState) { s.connClosing = true })
 	return c.wait(false)
 }
 
+// Closing returns a channel that is closed when Close is first called, i.e. when
+// the connection begins tearing down. Background best-effort work can select on
+// it to stop promptly rather than delaying teardown.
+func (c *Connection) Closing() <-chan struct{} {
+	return c.closing
+}
+
 // readIncoming collects inbound messages from the reader and delivers them, either responding
 // to outgoing calls or feeding requests to the queue.
 func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter Preempter) {

mcp/streamable_client_test.go

@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -2029,3 +2030,52 @@ func TestStreamableClient_StatelessSubscriptionsListen404(t *testing.T) {
 		t.Fatal("subscriptions/listen was not called")
 	}
 }
+
+// TestStreamableClientConnect_RespectsDeadline checks that Connect returns
+// close to its own deadline against an unresponsive peer (one that accepts the
+// TCP connection but never replies), rather than blocking on best-effort
+// cleanup. Regression test for #1189.
+func TestStreamableClientConnect_RespectsDeadline(t *testing.T) {
+	// Black hole: accept connections but never read or respond.
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	go func() {
+		var held []net.Conn
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				for _, h := range held {
+					h.Close()
+				}
+				return
+			}
+			held = append(held, c) // hold open, never respond
+		}
+	}()
+
+	const deadline = time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), deadline)
+	defer cancel()
+
+	client := NewClient(testImpl, nil)
+	start := time.Now()
+	_, err = client.Connect(ctx, &StreamableClientTransport{
+		Endpoint:   "http://" + ln.Addr().String(),
+		MaxRetries: 0,
+	}, nil)
+	elapsed := time.Since(start)
+
+	if err == nil {
+		t.Fatal("Connect unexpectedly succeeded against a black-hole server")
+	}
+	// The deadline bounds the initialize request; cleanup after it fails must
+	// not add another full notifyCancellationTimeout. Allow generous slack for
+	// slow CI while still catching the ~notifyCancellationTimeout overshoot.
+	if max := deadline + notifyCancellationTimeout - time.Second; elapsed > max {
+		t.Errorf("Connect returned after %v, want <= %v (deadline %v)",
+			elapsed.Round(10*time.Millisecond), max, deadline)
+	}
+}

mcp/transport.go

@@ -302,6 +302,17 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
 		go func() {
 			notifyCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout)
 			defer stop()
+			// If the connection starts tearing down, abandon the best-effort
+			// notify instead of blocking Close for up to notifyCancellationTimeout
+			// on an unresponsive peer: closing the connection already signals the
+			// cancellation. See issue #1189.
+			go func() {
+				select {
+				case <-conn.Closing():
+					stop()
+				case <-notifyCtx.Done():
+				}
+			}()
 			_ = conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{
 				Reason:    ctx.Err().Error(),
 				RequestID: call.ID().Raw(),