mcp: keep notifications/cancelled across session Close

docs/mcpgodebug.md

@@ -100,10 +100,11 @@ Options listed below were added and will be removed in the 1.9.0 version of the
   delivered (up to `notifyCancellationTimeout`, currently 5 s) before
   returning to the caller, restoring the previous behavior. The delivery
   error is joined into the caller's returned error. The default behavior was
-  changed so that the call is retired immediately and the notification is
-  sent asynchronously off the caller's return path: the caller returns as
-  soon as its context is cancelled and cannot be delayed by a slow or
-  unresponsive peer. See issue #1150.
+  changed so that the notification is sent asynchronously off the caller's
+  return path: the caller returns as soon as its context is cancelled and
+  cannot be delayed by a slow or unresponsive peer. The pending write is
+  reserved before return so a subsequent Close still delivers it. See
+  issues #1150 and #1505.
 
 Options below were removed, according to plan:
 

internal/docs/mcpgodebug.src.md

@@ -99,10 +99,11 @@ Options listed below were added and will be removed in the 1.9.0 version of the
   delivered (up to `notifyCancellationTimeout`, currently 5 s) before
   returning to the caller, restoring the previous behavior. The delivery
   error is joined into the caller's returned error. The default behavior was
-  changed so that the call is retired immediately and the notification is
-  sent asynchronously off the caller's return path: the caller returns as
-  soon as its context is cancelled and cannot be delayed by a slow or
-  unresponsive peer. See issue #1150.
+  changed so that the notification is sent asynchronously off the caller's
+  return path: the caller returns as soon as its context is cancelled and
+  cannot be delayed by a slow or unresponsive peer. The pending write is
+  reserved before return so a subsequent Close still delivers it. See
+  issues #1150 and #1505.
 
 Options below were removed, according to plan:
 

internal/jsonrpc2/conn.go

@@ -293,10 +293,12 @@ func (c *Connection) Notify(ctx context.Context, method string, params any) (err
 
 	c.updateInFlight(func(s *inFlightState) {
 		// If the connection is shutting down, allow outgoing notifications only if
-		// there is at least one call still in flight. The number of calls in flight
-		// cannot increase once shutdown begins, and allowing outgoing notifications
-		// may permit notifications that will cancel in-flight calls.
-		if len(s.outgoingCalls) == 0 && len(s.incomingByID) == 0 {
+		// there is at least one call still in flight, or a notification already
+		// reserved with [Connection.BeginOutgoingNotification]. The number of
+		// calls in flight cannot increase once shutdown begins, and allowing
+		// outgoing notifications may permit notifications that will cancel
+		// in-flight calls.
+		if len(s.outgoingCalls) == 0 && len(s.incomingByID) == 0 && s.outgoingNotifications == 0 {
 			err = s.shuttingDown(ErrConnClosing)
 			if err != nil {
 				return
@@ -317,6 +319,32 @@ func (c *Connection) Notify(ctx context.Context, method string, params any) (err
 	return c.write(ctx, notify)
 }
 
+// BeginOutgoingNotification records that a notification will be sent
+// asynchronously. Close waits until the returned function is called, and
+// [Connection.Notify] is allowed to write during shutdown while the
+// reservation is held. If the connection is already done, the reservation is
+// a no-op.
+func (c *Connection) BeginOutgoingNotification() (end func()) {
+	accepted := false
+	c.updateInFlight(func(s *inFlightState) {
+		select {
+		case <-c.done:
+			return
+		default:
+		}
+		s.outgoingNotifications++
+		accepted = true
+	})
+	if !accepted {
+		return func() {}
+	}
+	return sync.OnceFunc(func() {
+		c.updateInFlight(func(s *inFlightState) {
+			s.outgoingNotifications--
+		})
+	})
+}
+
 // Call invokes the target method and returns an object that can be used to await the response.
 // The params will be marshaled to JSON before sending over the wire, and will
 // be handed to the method invoked.

internal/jsonrpc2/conn_test.go

@@ -14,6 +14,7 @@ import (
 	"strings"
 	"sync"
 	"testing"
+	"time"
 )
 
 type fakeReader struct {
@@ -259,3 +260,61 @@ func (w *dropWriter) drops() []ID {
 	defer w.mu.Unlock()
 	return slices.Clone(w.dropped)
 }
+
+// TestBeginOutgoingNotificationAllowsNotifyDuringClose verifies that a
+// notification reserved with BeginOutgoingNotification is still written when
+// Close races with the sender, which is how MCP delivers notifications/cancelled
+// after a cancelled call (issue #1505).
+func TestBeginOutgoingNotificationAllowsNotifyDuringClose(t *testing.T) {
+	in := &chanReader{msgs: make(chan Message)}
+	w := fakeWriter{out: make(chan Message, 2)}
+	conn := NewConnection(context.Background(), ConnectionConfig{
+		Reader: in,
+		Writer: w,
+		Closer: in,
+		Bind: func(*Connection) Handler {
+			return HandlerFunc(func(context.Context, *Request) (any, error) {
+				return nil, ErrNotHandled
+			})
+		},
+	})
+
+	callCtx, cancel := context.WithCancel(context.Background())
+	call := conn.Call(callCtx, "slow", nil)
+	got := <-w.out
+	if req, ok := got.(*Request); !ok || !req.IsCall() {
+		t.Fatalf("got %#v, want the outgoing call", got)
+	}
+	cancel()
+	if err := call.Await(callCtx, nil); !errors.Is(err, context.Canceled) {
+		t.Fatalf("Await: %v, want context.Canceled", err)
+	}
+
+	end := conn.BeginOutgoingNotification()
+	conn.Retire(call, callCtx.Err())
+	go func() {
+		defer end()
+		_ = conn.Notify(context.Background(), "notifications/cancelled", nil)
+	}()
+
+	closed := make(chan error, 1)
+	go func() { closed <- conn.Close() }()
+
+	select {
+	case msg := <-w.out:
+		req, ok := msg.(*Request)
+		if !ok || req.IsCall() || req.Method != "notifications/cancelled" {
+			t.Fatalf("got %#v, want notifications/cancelled", msg)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("notification was not written")
+	}
+	select {
+	case err := <-closed:
+		if err != nil {
+			t.Errorf("Close: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Close did not return")
+	}
+}

mcp/streamable_test.go

@@ -4032,13 +4032,17 @@ func TestCallCancellation_FastReturn(t *testing.T) {
 	// slack accounts for scheduling jitter around the deadline.
 	const slack = 500 * time.Millisecond
 
+	stopStall := make(chan struct{})
 	stallCancelNotify := func(h http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if r.Method == http.MethodPost && r.Body != nil {
 				body, _ := io.ReadAll(r.Body)
 				r.Body = io.NopCloser(bytes.NewReader(body))
 				if bytes.Contains(body, []byte(notificationCancelled)) {
-					time.Sleep(stallDuration)
+					select {
+					case <-time.After(stallDuration):
+					case <-stopStall:
+					}
 				}
 			}
 			h.ServeHTTP(w, r)
@@ -4068,7 +4072,10 @@ func TestCallCancellation_FastReturn(t *testing.T) {
 	if err != nil {
 		t.Fatalf("client connect: %v", err)
 	}
-	t.Cleanup(func() { cs.Close() })
+	t.Cleanup(func() {
+		close(stopStall)
+		cs.Close()
+	})
 
 	callCtx, cancel := context.WithTimeout(context.Background(), callerDeadline)
 	defer cancel()
@@ -4087,6 +4094,60 @@ func TestCallCancellation_FastReturn(t *testing.T) {
 	}
 }
 
+// TestCloseAfterCanceledCallCancelsTheServerHandler verifies that closing the
+// session immediately after a cancelled tools/call still delivers
+// notifications/cancelled. Otherwise a stateful streamable handler keeps
+// running and ClientSession.Close blocks for closeDeleteTimeout. See #1505.
+func TestCloseAfterCanceledCallCancelsTheServerHandler(t *testing.T) {
+	entered := make(chan struct{}, 1)
+	canceled := make(chan struct{}, 1)
+	release := make(chan struct{})
+
+	server := NewServer(testImpl, nil)
+	AddTool(server, &Tool{Name: "wait"},
+		func(ctx context.Context, _ *CallToolRequest, _ struct{}) (*CallToolResult, any, error) {
+			entered <- struct{}{}
+			select {
+			case <-ctx.Done():
+				canceled <- struct{}{}
+				return nil, nil, ctx.Err()
+			case <-release:
+				return &CallToolResult{}, nil, nil
+			}
+		})
+	httpServer := httptest.NewServer(NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil))
+	t.Cleanup(httpServer.Close)
+	t.Cleanup(func() { close(release) })
+
+	client := NewClient(testImpl, nil)
+	session, err := client.Connect(context.Background(), &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		<-entered
+		cancel()
+	}()
+	if _, err := session.CallTool(ctx, &CallToolParams{Name: "wait"}); !errors.Is(err, context.Canceled) {
+		t.Fatalf("CallTool: got %v, want context.Canceled", err)
+	}
+
+	start := time.Now()
+	_ = session.Close()
+	closed := time.Since(start)
+
+	select {
+	case <-canceled:
+		if closed >= closeDeleteTimeout {
+			t.Errorf("Close took %v; want less than closeDeleteTimeout (%v)", closed, closeDeleteTimeout)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatalf("server handler never saw the cancellation; Close took %v", closed.Round(time.Millisecond))
+	}
+}
+
 // TestStreamableStateless_AcceptsNewProtocol is the positive control:
 // confirms that a stateless server still accepts new-protocol requests
 // (the rejection in TestStreamableStateful_RejectsNewProtocol must not

mcp/transport.go

@@ -32,8 +32,8 @@ const notifyCancellationTimeout = 5 * time.Second
 // blockingcancelnotify, when set to "1" via MCPGODEBUG, restores the previous
 // behavior of blocking the caller's return on delivery of the best-effort
 // notifications/cancelled message (up to notifyCancellationTimeout). By
-// default, the call is retired immediately and the notification is sent
-// asynchronously so the caller returns as soon as its context is cancelled.
+// default, the notification is sent asynchronously so the caller returns as
+// soon as its context is cancelled.
 var blockingcancelnotify = mcpgodebug.Value("blockingcancelnotify")
 
 // ErrConnectionClosed is returned when sending a message to a connection that
@@ -319,20 +319,26 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
 		}
 		return fmt.Errorf("%w: calling %q", ErrConnectionClosed, method)
 	case ctx.Err() != nil:
-		// The notifications/cancelled message is best-effort. Retire the call
-		// immediately (so an unresponsive peer cannot delay the eager
-		// retirement that cancelCall's docstring promises) and send the
-		// notification off the caller's return path so a slow or unresponsive
-		// peer cannot delay the caller past its own deadline. See issue #1150.
+		// The notifications/cancelled message is best-effort. Send it off the
+		// caller's return path so a slow or unresponsive peer cannot delay
+		// the caller past its own deadline (issue #1150).
+		//
+		// Reserve the write before returning so Close, which follows a
+		// cancelled call on graceful shutdown, waits for it and does not
+		// reject Notify. Retire immediately so an unresponsive peer's missing
+		// response cannot delay Close beyond notifyCancellationTimeout.
+		// See issue #1505.
 		//
 		// Setting MCPGODEBUG=blockingcancelnotify=1 restores the previous
 		// behavior of waiting synchronously for delivery inside cancelCall.
 		if blockingcancelnotify == "1" {
 			err := cancelCall(ctx, conn, call)
 			return errors.Join(ctx.Err(), err)
 		}
+		endNotify := conn.BeginOutgoingNotification()
 		conn.Retire(call, ctx.Err())
 		go func() {
+			defer endNotify()
 			notifyCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout)
 			defer stop()
 			_ = conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{