[undo] IJPL-201535 Ensure undo is valid if cancellation happens before command start

GitOrigin-RevId: 99bfc2b81f0ca0d41fa853d4ddaa87c8b74184d4

java/java-tests/testSrc/com/intellij/openapi/command/undo/UndoReliabilityTestKt.kt

@@ -1,14 +1,24 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.openapi.command.undo
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.EDT
 import com.intellij.openapi.application.runUndoTransparentWriteAction
+import com.intellij.openapi.command.CommandProcessor
+import com.intellij.openapi.command.impl.UndoManagerImpl
+import com.intellij.openapi.command.impl.UndoProvider
+import com.intellij.openapi.progress.Cancellation
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
 import com.intellij.testFramework.executeSomeCoroutineTasksAndDispatchAllInvocationEvents
-import kotlinx.coroutines.*
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
 import kotlin.coroutines.EmptyCoroutineContext
 
 
-class UndoReliabilityTestKt : ComplexUndoTest() {
+class UndoReliabilityTestKt : EditorUndoTestCase() {
 
   fun `test undo is valid if cancellation happens inside a command`() {
     newCoroutineScope().launch(Dispatchers.EDT) {
@@ -18,14 +28,84 @@ class UndoReliabilityTestKt : ComplexUndoTest() {
         firstEditor.document.insertString(1, "B")
       }
     }
-    executeSomeCoroutineTasksAndDispatchAllInvocationEvents(myProject)
+    executeSomeCoroutineTasksAndDispatchAllInvocationEvents()
     checkEditorText("AB")
     assertUndoInFirstEditorIsAvailable()
     undoFirstEditor()
     checkEditorText("")
   }
 
+  fun `test undo is valid if cancellation happens before command start`() {
+    forceCurrentEditorProviderService()
+    newCoroutineScope().launch(Dispatchers.EDT) {
+      val cancellationBeforeCommandStart = object : UndoProvider {
+        override fun commandStarted(project: Project?) = coroutineContext.cancel()
+        override fun commandFinished(project: Project?) = Unit
+      }
+      withUndoProvider(cancellationBeforeCommandStart) {
+        CommandProcessor.getInstance().runUndoTransparentAction {
+          // cancellation happens here inside the undo manager
+          runWriteActionWithoutCancellationCheckOnLockAcquiring {
+            firstEditor.document.insertString(0, "A")
+          }
+        }
+      }
+    }
+    executeSomeCoroutineTasksAndDispatchAllInvocationEvents()
+    checkEditorText("A")
+    assertUndoInFirstEditorIsAvailable()
+    undoFirstEditor()
+    checkEditorText("")
+  }
+
+  /**
+   * Force running write action within a cancelled coroutine
+   */
+  private fun runWriteActionWithoutCancellationCheckOnLockAcquiring(alreadyCanceledTask: () -> Unit) {
+    require(isCancelableSection())
+    val suppressCancellationCheckInWriteLockAcquiring = Cancellation.withNonCancelableSection()
+    var cancellationTokenClosed = false
+    try {
+      ApplicationManager.getApplication().runWriteAction {
+        suppressCancellationCheckInWriteLockAcquiring.finish()
+        cancellationTokenClosed = true
+        check(isCancelableSection()) {
+          "non cancelable section is not finished as expected"
+        }
+        alreadyCanceledTask.invoke()
+      }
+    } finally {
+      if (!cancellationTokenClosed) {
+        suppressCancellationCheckInWriteLockAcquiring.finish()
+      }
+    }
+  }
+
+  /**
+   * Force `CurrentEditorProvider.getInstance()` in order to match the production behavior
+   */
+  private fun forceCurrentEditorProviderService() {
+    (UndoManager.getGlobalInstance() as UndoManagerImpl).setOverriddenEditorProvider(null)
+    myManager.setOverriddenEditorProvider(null)
+  }
+
+  private fun withUndoProvider(undoProvider: UndoProvider, task: () -> Unit) {
+    val disposable = Disposer.newDisposable()
+    UndoProvider.EP_NAME.point.registerExtension(undoProvider, disposable)
+    UndoProvider.PROJECT_EP_NAME.getPoint(myProject).registerExtension(undoProvider, disposable)
+    try {
+      task.invoke()
+    } finally {
+      Disposer.dispose(disposable)
+    }
+  }
+
+  private fun isCancelableSection(): Boolean {
+    return !Cancellation.isInNonCancelableSection()
+  }
+
   private fun newCoroutineScope(): CoroutineScope {
+    @Suppress("RAW_SCOPE_CREATION")
     return CoroutineScope(EmptyCoroutineContext)
   }
 }

platform/core-impl/src/com/intellij/openapi/command/impl/CommandIdService.java

@@ -3,6 +3,7 @@ package com.intellij.openapi.command.impl;
 
 import com.intellij.openapi.application.Application;
 import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.progress.ProgressManager;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -15,7 +16,9 @@ public interface CommandIdService {
   static @Nullable CommandIdService getInstance() {
     Application application = ApplicationManager.getApplication();
     if (application != null) {
-      CommandIdService service = application.getService(CommandIdService.class);
+      CommandIdService service = ProgressManager.getInstance().computeInNonCancelableSection(
+        () -> application.getService(CommandIdService.class)
+      );
       return service;
     }
     return null;

platform/platform-impl/src/com/intellij/openapi/command/impl/CmdEvent.java

@@ -21,6 +21,9 @@ public interface CmdEvent {
   boolean isTransparent();
 
   default @NotNull CmdEvent withProject(@Nullable Project project) {
+    if (project == project()) {
+      return this;
+    }
     return create(
       id(),
       project,
@@ -33,6 +36,9 @@ public interface CmdEvent {
   }
 
   default @NotNull CmdEvent withRecordOriginalDocument(boolean recordOriginalDocument) {
+    if (recordOriginalDocument == recordOriginalDocument()) {
+      return this;
+    }
     return create(
       id(),
       project(),

platform/platform-impl/src/com/intellij/openapi/command/impl/CommandBuilder.java

@@ -61,7 +61,7 @@ final class CommandBuilder {
   }
 
   void commandStarted(@NotNull CmdEvent cmdEvent, @NotNull CurrentEditorProvider editorProvider) {
-    assertOutsideCommand();
+    assertOutsideCommand(cmdEvent);
     if (LOG.isTraceEnabled() || ApplicationManager.getApplication().isUnitTestMode()) {
       this.tracedStartCommand = new Throwable();
     }
@@ -113,7 +113,7 @@ final class CommandBuilder {
   }
 
   @NotNull PerformedCommand commandFinished(@NotNull CmdEvent cmdEvent) {
-    assertInsideCommand();
+    assertInsideCommand(cmdEvent);
     if (isGroupIdChangeSupported) {
       this.cmdEvent = cmdEvent;
     }
@@ -124,16 +124,27 @@ final class CommandBuilder {
     return buildAndReset();
   }
 
+  void assertOutsideCommand() {
+    assertOutsideCommand(null);
+  }
 
-  void assertInsideCommand() {
-    if (!isInsideCommand) {
-      throw new UndoIllegalStateException("Must be called inside a command");
+  private void assertInsideCommand() {
+    assertInsideCommand(null);
+  }
+
+  private void assertOutsideCommand(@Nullable CmdEvent cmdEvent) {
+    if (isInsideCommand) {
+      String startEvent = cmdEvent == null ? "" : (", startEvent: " + cmdEvent);
+      throw new UndoIllegalStateException(
+        "Nested command detected, please report the stacktrace" + startEvent, tracedStartCommand
+      );
     }
   }
 
-  void assertOutsideCommand() {
-    if (isInsideCommand) {
-      throw new UndoIllegalStateException("Nested command detected, please report the stacktrace", tracedStartCommand);
+  private void assertInsideCommand(@Nullable CmdEvent cmdEvent) {
+    if (!isInsideCommand) {
+      String finishEvent = cmdEvent == null ? "" : (", finishEvent: " + cmdEvent);
+      throw new UndoIllegalStateException("Must be called inside a command" + finishEvent);
     }
   }
 

platform/platform-impl/src/com/intellij/openapi/command/impl/StableEditorProvider.java

@@ -3,6 +3,7 @@ package com.intellij.openapi.command.impl;
 
 import com.intellij.openapi.fileEditor.FileEditor;
 import com.intellij.openapi.fileEditor.impl.CurrentEditorProvider;
+import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.Project;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -22,7 +23,9 @@ final class StableEditorProvider implements CurrentEditorProvider {
     if (isInitialized) {
       return editor;
     }
-    editor = provider.getCurrentEditor(project);
+    editor = ProgressManager.getInstance().computeInNonCancelableSection(
+      () -> provider.getCurrentEditor(project)
+    );
     isInitialized = true;
     return editor;
   }

platform/platform-impl/src/com/intellij/openapi/command/impl/UndoManagerImpl.java

@@ -205,8 +205,11 @@ public class UndoManagerImpl extends UndoManager {
   }
 
   public @NotNull CurrentEditorProvider getEditorProvider() {
-    CurrentEditorProvider provider = myOverriddenEditorProvider;
-    return new StableEditorProvider(provider != null ? provider : CurrentEditorProvider.getInstance());
+    CurrentEditorProvider overriddenProvider = myOverriddenEditorProvider;
+    CurrentEditorProvider editorProvider = overriddenProvider != null
+        ? overriddenProvider
+        : ProgressManager.getInstance().computeInNonCancelableSection(CurrentEditorProvider::getInstance);
+    return new StableEditorProvider(editorProvider);
   }
 
   public @Nullable Project getProject() {
@@ -531,9 +534,11 @@ public class UndoManagerImpl extends UndoManager {
   }
 
   private @NotNull List<UndoProvider> getUndoProviders() {
-    return myProject == null
-           ? UndoProvider.EP_NAME.getExtensionList()
-           : UndoProvider.PROJECT_EP_NAME.getExtensionList(myProject);
+    return ProgressManager.getInstance().computeInNonCancelableSection(() -> {
+      return myProject == null
+      ? UndoProvider.EP_NAME.getExtensionList()
+      : UndoProvider.PROJECT_EP_NAME.getExtensionList(myProject);
+    });
   }
 
   private @NotNull ComponentManager getComponentManager() {