From b4d1381fe740258ba27e88892a1c5244c392bb42 Mon Sep 17 00:00:00 2001 From: Arpan Mondal Date: Thu, 3 Sep 2026 12:09:00 +0530 Subject: [PATCH] mcp: don't let a best-effort cancel notify delay connection teardown When a call is cancelled at its deadline, the notifications/cancelled is sent asynchronously, but Connection.Close's wait still blocked on that in-flight write for up to notifyCancellationTimeout against an unresponsive peer. Add a Closing signal to the jsonrpc2 connection and abort the notify once teardown begins, so Connect returns within its own deadline. Fixes #1189 --- internal/jsonrpc2/conn.go | 15 +++++++++++ mcp/streamable_client_test.go | 50 +++++++++++++++++++++++++++++++++++ mcp/transport.go | 11 ++++++++ 3 files changed, 76 insertions(+) diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 5849a871..683a1857 100644 --- a/internal/jsonrpc2/conn.go +++ b/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) { diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 77ff81a2..3a067e22 100644 --- a/mcp/streamable_client_test.go +++ b/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) + } +} diff --git a/mcp/transport.go b/mcp/transport.go index 1cf03251..c0e1ae89 100644 --- a/mcp/transport.go +++ b/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(),