[analysis_server] Handle request cancellation when a request is waiting for an outbound user prompt

This updates the `showUserPrompt` method to accept a cancellation token so that if requests are cancelled (by the client, or something like a second refactor cancelling the first), the prompt is also cancelled (and when the response arrives, it can be ignored).

Note: Unfortunately the protocol (and VS Code) don't actually allow for the prompt to be cancelled/hidden from the user, this is mainly to avoid us keeping the request "alive" on the server if the user ignores a prompt and we know the parent request was cancelled anyway.

Fixes https://github.com/dart-lang/sdk/issues/63285

Change-Id: Ia77880a749b284a2ff31c8cf560307ad908d2175
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/499360
Reviewed-by: Konstantin Shcheglov <scheglov@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Brian Wilkerson <brianwilkerson@google.com>

pkg/analysis_server/lib/src/analysis_server.dart

@@ -84,6 +84,7 @@ import 'package:analyzer/src/generated/sdk.dart';
 import 'package:analyzer/src/util/file_paths.dart' as file_paths;
 import 'package:analyzer/src/util/performance/operation_performance.dart';
 import 'package:analyzer/src/util/platform_info.dart';
+import 'package:analyzer/src/utilities/cancellation.dart';
 import 'package:analyzer/src/utilities/extensions/analysis_session.dart';
 import 'package:analyzer/src/workspace/basic.dart';
 import 'package:analyzer/src/workspace/blaze.dart';
@@ -108,6 +109,7 @@ typedef UserPromptSender =
       MessageType type,
       String message,
       List<String> actionLabels,
+      lsp.CancellationToken cancellationToken,
     );
 
 /// Implementations of [AnalysisServer] implement a server that listens
@@ -556,7 +558,9 @@ abstract class AnalysisServer {
     }
 
     unawaited(
-      prompt(MessageType.info, unifiedAnalytics.getConsentMessage, ['Ok']),
+      prompt(MessageType.info, unifiedAnalytics.getConsentMessage, [
+        lsp.UserPromptActions.ok,
+      ], NotCancelableToken()),
     );
     unifiedAnalytics.clientShowedMessage();
   }
@@ -1140,6 +1144,7 @@ abstract class AnalysisServer {
     MessageType type,
     String message,
     List<String> actionLabels,
+    lsp.CancellationToken cancellationToken,
   );
 
   @mustCallSuper

pkg/analysis_server/lib/src/legacy_analysis_server.dart

@@ -978,6 +978,7 @@ class LegacyAnalysisServer extends AnalysisServer {
     MessageType type,
     String message,
     List<String> actionLabels,
+    CancellationToken cancellationToken,
   ) async {
     assert(supportsShowMessageRequest);
     var requestId = (nextServerRequestId++).toString();
@@ -987,8 +988,18 @@ class LegacyAnalysisServer extends AnalysisServer {
       message,
       actions,
     ).toRequest(requestId, clientUriConverter: uriConverter);
-    var response = await sendRequest(request);
-    return response.result?['action'] as String?;
+    var responseFuture = sendRequest(request);
+
+    // Wait for either the result, or cancellation.
+    await Future.any([responseFuture, cancellationToken.whenCancelled]);
+
+    if (cancellationToken.isCancellationRequested) {
+      return null;
+    } else {
+      // If we didn't enter the branch above, we know this future completed.
+      var response = await responseFuture;
+      return response.result?['action'] as String?;
+    }
   }
 
   @override

pkg/analysis_server/lib/src/lsp/constants.dart

@@ -396,6 +396,7 @@ abstract final class ServerErrorCodes {
 abstract final class UserPromptActions {
   static const String yes = 'Yes';
   static const String no = 'No';
+  static const String ok = 'Ok';
   static const String cancel = 'Cancel';
   static const String renameAnyway = 'Rename Anyway';
 }

pkg/analysis_server/lib/src/lsp/handlers/handler_rename.dart

@@ -205,9 +205,9 @@ class RenameHandler extends LspMessageHandler<RenameParams, WorkspaceEdit?> {
           return error(ServerErrorCodes.renameNotValid, finalStatus.message!);
         }
 
-        // Set the completer to complete to show that request is paused, and
-        // that processing of incoming messages can continue while we wait
-        // for the user's response.
+        // Complete the message before we make the outbound request because when
+        // the server is in non-overlapping request mode, we cannot have the
+        // server stall because this request is blocked on user-input.
         message.completer?.complete();
 
         // Otherwise, ask the user whether to proceed with the rename.
@@ -215,6 +215,7 @@ class RenameHandler extends LspMessageHandler<RenameParams, WorkspaceEdit?> {
           MessageType.warning,
           finalStatus.message!,
           [UserPromptActions.renameAnyway, UserPromptActions.cancel],
+          token,
         );
 
         if (token.isCancellationRequested) {
@@ -278,8 +279,12 @@ class RenameHandler extends LspMessageHandler<RenameParams, WorkspaceEdit?> {
             var shouldRename =
                 renameConfig == 'always' ||
                 (renameConfig == 'prompt' &&
-                    await _promptToRenameFile(actualFilename, newFilename));
-            if (shouldRename) {
+                    await _promptToRenameFile(
+                      actualFilename,
+                      newFilename,
+                      token,
+                    ));
+            if (shouldRename && !token.isCancellationRequested) {
               var newPath = pathContext.join(folder, newFilename);
               var renameEdit = createRenameEdit(
                 uriConverter,
@@ -292,6 +297,10 @@ class RenameHandler extends LspMessageHandler<RenameParams, WorkspaceEdit?> {
         }
       }
 
+      if (token.isCancellationRequested) {
+        return cancelled(token);
+      }
+
       return success(workspaceEdit);
     });
   }
@@ -311,6 +320,7 @@ class RenameHandler extends LspMessageHandler<RenameParams, WorkspaceEdit?> {
   Future<bool> _promptToRenameFile(
     String oldFilename,
     String newFilename,
+    CancellationToken cancellationToken,
   ) async {
     var prompt = server.userPromptSender;
     // If we can't prompt, do the same as if they said no.
@@ -322,6 +332,7 @@ class RenameHandler extends LspMessageHandler<RenameParams, WorkspaceEdit?> {
       MessageType.info,
       "Rename '$oldFilename' to '$newFilename'?",
       [UserPromptActions.yes, UserPromptActions.no],
+      cancellationToken,
     );
 
     return userChoice == UserPromptActions.yes;

pkg/analysis_server/lib/src/lsp/lsp_analysis_server.dart

@@ -936,12 +936,14 @@ class LspAnalysisServer extends AnalysisServer {
     MessageType type,
     String message,
     List<String> actions,
+    CancellationToken cancellationToken,
   ) async {
     assert(supportsShowMessageRequest);
     var response = await showUserPromptItems(
       type,
       message,
       actions.map((title) => MessageActionItem(title: title)).toList(),
+      cancellationToken,
     );
     return response?.title;
   }
@@ -959,9 +961,10 @@ class LspAnalysisServer extends AnalysisServer {
     MessageType type,
     String message,
     List<MessageActionItem> actions,
+    CancellationToken cancellationToken,
   ) async {
     assert(supportsShowMessageRequest);
-    var response = await sendLspRequest(
+    var responseFuture = sendLspRequest(
       Method.window_showMessageRequest,
       ShowMessageRequestParams(
         type: type.forLsp,
@@ -970,10 +973,19 @@ class LspAnalysisServer extends AnalysisServer {
       ),
     );
 
-    var result = response.result;
-    return result != null
-        ? MessageActionItem.fromJson(response.result as Map<String, Object?>)
-        : null;
+    // Wait for either the result, or cancellation.
+    await Future.any([responseFuture, cancellationToken.whenCancelled]);
+
+    if (cancellationToken.isCancellationRequested) {
+      return null;
+    } else {
+      // If we didn't enter the branch above, we know this future completed.
+      var response = await responseFuture;
+      var result = response.result;
+      return result != null
+          ? MessageActionItem.fromJson(response.result as Map<String, Object?>)
+          : null;
+    }
   }
 
   @override

pkg/analysis_server/lib/src/scheduler/message_scheduler.dart

@@ -144,7 +144,15 @@ final class MessageScheduler {
           listener?.addActiveMessage(message);
           var id = _processCancellation(msg);
           listener?.messageCompleted(message, id: id);
-          return;
+          // Only skip adding to the queue if we processed the cancellation, as
+          // there are cases where an active message might not be in our queue
+          // because we sometimes pretend a request is completed when it isn't,
+          // but it has made an outbound reverse-request and we cannot stall
+          // the server while that's open if the server is in non-overlapping
+          // request mode.
+          if (id != null) {
+            return;
+          }
         } else if (method == lsp.Method.textDocument_didChange) {
           // Document change notifications are _not_ handled immediately, but
           // some active or pending requests can be cancelled before the normal

pkg/analysis_server/lib/src/services/user_prompts/dart_fix_prompt_manager.dart

@@ -155,6 +155,7 @@ class DartFixPromptManager {
   }) async {
     _hasPromptedThisSession = true;
 
+    var cancellationToken = NotCancelableToken(); // Can't be cancelled by user
     var executeCommandHandler = server.executeCommandHandler;
     String prompt;
     List<String> actions;
@@ -179,6 +180,7 @@ class DartFixPromptManager {
       MessageType.info,
       prompt,
       actions,
+      cancellationToken,
     ).then((value) => value, onError: (_) => null);
 
     switch ((response, executeCommandHandler)) {
@@ -196,6 +198,7 @@ class DartFixPromptManager {
             execHandler,
             userPromptSender,
             command,
+            cancellationToken,
           ),
         );
 
@@ -244,6 +247,7 @@ class DartFixPromptManager {
     ExecuteCommandHandler handler,
     UserPromptSender userPromptSender,
     String command,
+    CancellationToken cancellationToken,
   ) async {
     // Go through the main handle method so that things like analytics are
     // recorded the same.
@@ -256,7 +260,7 @@ class DartFixPromptManager {
         clientCapabilities: clientCapabilities,
         isTrustedCaller: true,
       ),
-      NotCancelableToken(),
+      cancellationToken,
     );
 
     result.ifError((error) {
@@ -265,6 +269,7 @@ class DartFixPromptManager {
           MessageType.error,
           "Failed to execute '$command': ${error.message}",
           [],
+          cancellationToken,
         ),
       );
     });

pkg/analysis_server/lib/src/services/user_prompts/survey_manager.dart

@@ -6,6 +6,7 @@ import 'dart:async';
 
 import 'package:analysis_server/src/analysis_server.dart';
 import 'package:analyzer/instrumentation/service.dart';
+import 'package:analyzer/src/utilities/cancellation.dart';
 import 'package:unified_analytics/unified_analytics.dart';
 
 /// An interface for interacting with surveys via the unified_analytics package.
@@ -80,6 +81,7 @@ class SurveyManager {
         MessageType.info,
         survey.description,
         buttonMap.keys.toList(),
+        NotCancelableToken(), // Not user-cancellable
       );
       var clickedButton = buttonMap[clickedButtonText];
       if (clickedButton == null) return;

pkg/analysis_server/test/domain_server_test.dart

@@ -11,6 +11,7 @@ import 'package:analysis_server/protocol/protocol_generated.dart'
     hide MessageType;
 import 'package:analysis_server/src/analysis_server.dart' show MessageType;
 import 'package:analysis_server/src/services/user_prompts/dart_fix_prompt_manager.dart';
+import 'package:analyzer/src/utilities/cancellation.dart';
 import 'package:test/test.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
 
@@ -378,6 +379,7 @@ class ServerDomainTest extends PubPackageAnalysisServerTest {
       MessageType.warning,
       'message',
       ['a', 'b'],
+      NotCancelableToken(),
     );
     expect(serverChannel.serverRequestsSent, hasLength(1));
 
@@ -400,6 +402,7 @@ class ServerDomainTest extends PubPackageAnalysisServerTest {
       MessageType.warning,
       'message',
       ['a', 'b'],
+      NotCancelableToken(),
     );
     expect(serverChannel.serverRequestsSent, hasLength(1));
 

pkg/analysis_server/test/lsp/code_actions_source_test.dart

@@ -169,7 +169,7 @@ void f() {
 
         return commandFuture;
       },
-      handler: (edit) {
+      handler: (edit) async {
         // When the server sends the edit back, just keep a copy and say we
         // applied successfully (we'll verify the actual edit below).
         editParams = edit;

pkg/analysis_server/test/lsp/rename_test.dart

@@ -2,6 +2,8 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'dart:async';
+
 import 'package:analysis_server/lsp_protocol/protocol.dart';
 import 'package:analysis_server/src/lsp/constants.dart';
 import 'package:analyzer/src/test_utilities/platform.dart';
@@ -22,6 +24,24 @@ void main() {
 
 @reflectiveTest
 class RenameTest extends AbstractLspAnalysisServerTest {
+  /// A completer that completes when an outbound rename request is sent.
+  ///
+  /// This is used in cancellation tests that need to know when the request has
+  /// been sent and access the ID to cancel.
+  Completer<RequestMessage> outboundRenameRequestCompleter = Completer();
+
+  @override
+  RequestMessage makeRenameRequest(
+    int? version,
+    Uri uri,
+    Position pos,
+    String newName,
+  ) {
+    var request = super.makeRenameRequest(version, uri, pos, newName);
+    outboundRenameRequestCompleter.complete(request);
+    return request;
+  }
+
   Future<void> test_prepare_class() {
     const content = '''
 class MyClass {}
@@ -548,6 +568,55 @@ final a = new MyOtherClass();
     verifyEdit(result, expectedContent);
   }
 
+  Future<void> test_rename_duplicateName_cancelWhilePromptPending() async {
+    const content = '''
+class MyOtherClass {}
+class MyClass {}
+final a = n^ew MyClass();
+''';
+    // Pass a completer as beforeResponding so we can hold off responding to
+    // the prompt reverse-request.
+    var reverseRequestStartedCompleter = Completer<void>();
+    var reverseRequestCompleter = Completer<void>();
+    var responseFuture = _test_rename_prompt(
+      content,
+      'MyOtherClass',
+      expectedMessage:
+          'Library already declares class with name \'MyOtherClass\'.',
+      // When we complete, we will send renameAnyway, but since we
+      // cancelled the request in the meantime, we expect a cancellation
+      // result.
+      action: UserPromptActions.renameAnyway,
+      beforeResponding: () {
+        // Mark that we started, so the test can continue.
+        reverseRequestStartedCompleter.complete();
+        // But wait for us to trigger the end.
+        return reverseRequestCompleter.future;
+      },
+    );
+
+    // Wait for the server to send the reverse-request so we don't cancel too
+    // early.
+    await reverseRequestStartedCompleter.future;
+
+    // Now, cancel the request.
+    await sendNotificationToServer(
+      makeNotification(
+        Method.cancelRequest,
+        CancelParams(id: (await outboundRenameRequestCompleter.future).id),
+      ),
+    );
+
+    // Expect that the response completes and returns a cancelled state.
+    var response = await responseFuture;
+    expect(response.result, isNull);
+    expect(response.error, isNotNull);
+    expect(response.error, isResponseError(ErrorCodes.RequestCancelled));
+
+    // Unblock the reverse-request.
+    reverseRequestCompleter.complete();
+  }
+
   Future<void> test_rename_duplicateName_hover_beforeResponding() async {
     const content = '''
 class MyOtherClass {}

pkg/analysis_server/test/lsp/request_helpers_mixin.dart

@@ -1228,15 +1228,16 @@ mixin LspReverseRequestHelpersMixin {
     Method method,
     R Function(Map<String, dynamic>) fromJson,
     Future<T> Function() f, {
-    required FutureOr<RR> Function(R) handler,
+    required Future<RR> Function(R) handler,
     Duration timeout = const Duration(seconds: 5),
-  }) async {
+  }) {
     late Future<T> outboundRequest;
     Object? outboundRequestError;
 
-    // Run [f] and wait for the incoming request from the server.
-    var incomingRequest =
-        await expectRequest(method, () {
+    // Execute [f] to start the outbound request that will trigger the inbound
+    // request.
+    var incomingRequestFuture =
+        expectRequest(method, () {
           // Don't return/await the response yet, as this may not complete until
           // after we have handled the request that comes from the server.
           outboundRequest = f();
@@ -1262,11 +1263,19 @@ mixin LspReverseRequestHelpersMixin {
           throw outboundRequestError ?? timeoutException;
         }, test: (e) => e is TimeoutException);
 
-    // Handle the request from the server and send the response back.
-    var clientsResponse = await handler(
-      fromJson(incomingRequest.params as Map<String, Object?>),
+    // When the inbound request arrives, send the response back. Don't wait for
+    // it here because some tests handle requests that complete early (eg.
+    // cancellation) and need to just await the outbound request.
+    unawaited(
+      incomingRequestFuture.then((incomingRequest) async {
+        // Call the handler to compute the repsonse.
+        var clientsResponse = await handler(
+          fromJson(incomingRequest.params as Map<String, Object?>),
+        );
+        // Send the repsonse back to the server.
+        respondTo(incomingRequest, clientsResponse);
+      }),
     );
-    respondTo(incomingRequest, clientsResponse);
 
     // Return a future that completes when the response to the original request
     // (from [f]) returns.
@@ -1349,7 +1358,7 @@ mixin LspVerifyEditHelpersMixin
           Method.workspace_applyEdit,
           ApplyWorkspaceEditParams.fromJson,
           function,
-          handler: (edit) {
+          handler: (edit) async {
             // When the server sends the edit back, just keep a copy and say we
             // applied successfully (it'll be verified by the caller).
             editParams = edit;
@@ -1377,7 +1386,7 @@ mixin LspVerifyEditHelpersMixin
     Method method,
     R Function(Map<String, dynamic>) fromJson,
     Future<T> Function() f, {
-    required FutureOr<RR> Function(R) handler,
+    required Future<RR> Function(R) handler,
     Duration timeout = const Duration(seconds: 5),
   });
 

pkg/analysis_server/test/lsp/server_abstract.dart

@@ -1256,7 +1256,7 @@ mixin LspAnalysisServerTestMixin
       Method.client_registerCapability,
       RegistrationParams.fromJson,
       f,
-      handler: (registrationParams) {
+      handler: (registrationParams) async {
         registrations.addAll(registrationParams.registrations);
       },
     );
@@ -1281,7 +1281,7 @@ mixin LspAnalysisServerTestMixin
       Method.client_unregisterCapability,
       UnregistrationParams.fromJson,
       f,
-      handler: (unregistrationParams) {
+      handler: (unregistrationParams) async {
         registrations.removeWhere(
           (element) => unregistrationParams.unregisterations.any(
             (u) => u.id == element.id,

pkg/analysis_server/test/services/user_prompts/dart_fix_prompt_manager_test.dart

@@ -395,7 +395,7 @@ class TestServer implements LspAnalysisServer {
 
   @override
   UserPromptSender? get userPromptSender => supportsShowMessageRequest
-      ? (_, promptText, promptActions) async {
+      ? (_, promptText, promptActions, cancellationToken) async {
           lastPromptText = promptText;
           lastPromptActions = promptActions;
           assert(promptActions.contains(respondToPromptWithAction));

pkg/analysis_server/test/services/user_prompts/survey_manager_test.dart

@@ -5,6 +5,7 @@
 import 'package:analysis_server/src/analysis_server.dart';
 import 'package:analysis_server/src/services/user_prompts/survey_manager.dart';
 import 'package:analyzer/instrumentation/instrumentation.dart';
+import 'package:analyzer/src/utilities/cancellation.dart';
 import 'package:analyzer_testing/resource_provider_mixin.dart';
 import 'package:test/test.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
@@ -217,6 +218,7 @@ class TestServer implements AnalysisServer {
     MessageType type,
     String message,
     List<String> actionLabels,
+    CancellationToken cancellationToken,
   ) async {
     return actionLabels.where((s) => s == respondWithButton).firstOrNull;
   }

pkg/analysis_server/test/shared/shared_code_actions_source_tests.dart

@@ -251,7 +251,7 @@ String? a;
           // Claim that we failed tpo apply the edits. This is what the client
           // would do if the edits provided were for an old version of the
           // document.
-          handler: (edit) => ApplyWorkspaceEditResult(
+          handler: (edit) async => ApplyWorkspaceEditResult(
             applied: false,
             failureReason: 'Document changed',
           ),

pkg/analyzer/lib/src/utilities/cancellation.dart

@@ -2,8 +2,10 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'dart:async';
+
 class CancelableToken extends CancellationToken {
-  bool _isCancelled = false;
+  final Completer<void> _cancelCompleter = Completer<void>();
   int? _cancellationCode;
   String? _cancellationReason;
 
@@ -14,10 +16,15 @@ class CancelableToken extends CancellationToken {
   String? get cancellationReason => _cancellationReason;
 
   @override
-  bool get isCancellationRequested => _isCancelled;
+  bool get isCancellationRequested => _cancelCompleter.isCompleted;
+
+  @override
+  Future<void> get whenCancelled => _cancelCompleter.future;
 
   void cancel({int? code, String? reason}) {
-    _isCancelled = true;
+    if (!_cancelCompleter.isCompleted) {
+      _cancelCompleter.complete();
+    }
     _cancellationCode = code;
     _cancellationReason = reason;
   }
@@ -27,11 +34,18 @@ class CancelableToken extends CancellationToken {
 /// to be skipped when a caller is no longer interested in the result, for example
 /// when a $/cancel request is received for an in-progress request.
 abstract class CancellationToken {
+  /// Whether cancellation has been requested.
   bool get isCancellationRequested;
+
+  /// A [Future] that completes if/when cancellation is requested.
+  Future<void> get whenCancelled;
 }
 
 /// A [CancellationToken] that cannot be cancelled.
 class NotCancelableToken extends CancellationToken {
+  @override
+  final Future<void> whenCancelled = Completer<void>().future;
+
   @override
   bool get isCancellationRequested => false;
 }