fix(test): fix batch test generation cancel button

Fixes issue where cancel button in batch test generation progress dialog was not responsive. Added proper cancellation checks throughout the test generation process and replaced executor-based approach with cancellable background tasks.

Resolves #407

Author: phodal

core/src/main/kotlin/cc/unitmesh/devti/actions/chat/AutoTestInMenuAction.kt

@@ -8,8 +8,7 @@ import cc.unitmesh.devti.intentions.action.test.TestCodeGenRequest
 import com.intellij.openapi.actionSystem.*
 import com.intellij.openapi.actionSystem.ActionPlaces.PROJECT_VIEW_POPUP
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.progress.EmptyProgressIndicator
-import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.*
 import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.MessageType
@@ -21,9 +20,11 @@ import com.intellij.openapi.wm.WindowManager
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.psi.impl.file.PsiDirectoryFactory
-import java.util.concurrent.Executors
+import com.intellij.openapi.diagnostic.logger
 
 class AutoTestInMenuAction : AnAction(AutoDevBundle.message("intentions.chat.code.test.name")) {
+    private val logger = logger<AutoTestInMenuAction>()
+
     override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.BGT
 
     fun getActionType(): ChatActionType = ChatActionType.GENERATE_TEST
@@ -60,26 +61,49 @@ class AutoTestInMenuAction : AnAction(AutoDevBundle.message("intentions.chat.cod
     }
 
     private fun batchGenerateTests(files: List<PsiFile>, project: Project, editor: Editor?) {
-        val total = files.size
-        val executor = Executors.newSingleThreadExecutor()
-        files.forEachIndexed { index, file ->
-            val task = TestCodeGenTask(
-                TestCodeGenRequest(file, file, project, editor),
-                AutoDevBundle.message("intentions.chat.code.test.name")
-            )
-
-            executor.submit {
-                val progressMessage = """${index + 1}/${total} Processing file ${file.name} for test generation"""
-                ProgressManager.getInstance().runProcessWithProgressSynchronously(
-                    {
-                        task.run(object : EmptyProgressIndicator() {})
-                    },
-                    progressMessage, true, project
-                )
+        val batchTask = object : Task.Backgroundable(
+            project,
+            AutoDevBundle.message("intentions.chat.code.test.name") + " (Batch)",
+            true
+        ) {
+            override fun run(indicator: ProgressIndicator) {
+                val total = files.size
+                indicator.isIndeterminate = false
+                indicator.fraction = 0.0
+
+                files.forEachIndexed { index, file ->
+                    // Check for cancellation before processing each file
+                    indicator.checkCanceled()
+
+                    indicator.text = "Processing ${index + 1}/$total: ${file.name}"
+                    indicator.fraction = index.toDouble() / total
+
+                    val task = TestCodeGenTask(
+                        TestCodeGenRequest(file, file, project, editor),
+                        AutoDevBundle.message("intentions.chat.code.test.name")
+                    )
+
+                    try {
+                        task.run(indicator)
+                        indicator.fraction = (index + 1).toDouble() / total
+                    } catch (e: ProcessCanceledException) {
+                        // User cancelled, stop processing
+                        indicator.text = "Batch test generation cancelled"
+                        throw e
+                    } catch (e: Exception) {
+                        // Log error but continue with next file
+                        logger.warn("Failed to generate test for file: ${file.name}", e)
+                        indicator.fraction = (index + 1).toDouble() / total
+                    }
+                }
+
+                indicator.fraction = 1.0
+                indicator.text = "Batch test generation completed"
             }
         }
 
-        executor.shutdown()
+        ProgressManager.getInstance()
+            .runProcessWithProgressAsynchronously(batchTask, BackgroundableProcessIndicator(batchTask))
     }
 
     private fun isEnabled(e: AnActionEvent): Boolean {

core/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/TestCodeGenTask.kt

@@ -26,6 +26,7 @@ import com.intellij.openapi.editor.ScrollType
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.progress.ProgressIndicator
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.progress.Task
 import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.Project
@@ -53,10 +54,16 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
         indicator.fraction = 0.1
         indicator.text = AutoDevBundle.message("intentions.chat.code.test.step.prepare-context")
 
+        // Check for cancellation early
+        indicator.checkCanceled()
+
         AutoDevStatusService.notifyApplication(AutoDevStatus.InProgress)
         val testContext = autoTestService.findOrCreateTestFile(request.file, request.project, request.element)
         DumbService.getInstance(request.project).waitForSmartMode()
 
+        // Check for cancellation after waiting for smart mode
+        indicator.checkCanceled()
+
         if (testContext == null) {
             AutoDevStatusService.notifyApplication(AutoDevStatus.Error)
             logger.error("Failed to create test file for: ${request.file}")
@@ -66,16 +73,27 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
         indicator.text = AutoDevBundle.message("intentions.chat.code.test.step.collect-context")
         indicator.fraction = 0.3
 
+        // Check for cancellation before collecting context
+        indicator.checkCanceled()
+
         val testPromptContext = TestCodeGenContext()
 
         val creationContext =
             ChatCreationContext(ChatOrigin.Intention, actionType, request.file, listOf(), element = request.element)
 
         val contextItems: List<ChatContextItem> = runBlocking {
+            // Check for cancellation in the blocking context
+            if (indicator.isCanceled) {
+                throw ProcessCanceledException()
+            }
             ChatContextProvider.collectChatContextList(request.project, creationContext)
         }
 
         testPromptContext.frameworkContext = contextItems.joinToString("\n", transform = ChatContextItem::text)
+
+        // Check for cancellation before read actions
+        indicator.checkCanceled()
+
         ReadAction.compute<Unit, Throwable> {
             if (testContext.relatedClasses.isNotEmpty()) {
                 testPromptContext.relatedClasses = testContext.relatedClasses.joinToString("\n") {
@@ -113,6 +131,9 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
         testPromptContext.isNewFile = testContext.isNewFile
         testPromptContext.extContext = getCustomAgentTestContext(testPromptContext)
 
+        // Check for cancellation before template rendering
+        indicator.checkCanceled()
+
         templateRender.context = testPromptContext
         val prompter = templateRender.renderTemplate(template)
 
@@ -121,6 +142,9 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
         indicator.fraction = 0.6
         indicator.text = AutoDevBundle.message("intentions.request.background.process.title")
 
+        // Check for cancellation before LLM request
+        indicator.checkCanceled()
+
         val flow: Flow<String> = try {
             LlmFactory.create(request.project).stream(prompter, "", false)
         } catch (e: Exception) {
@@ -130,13 +154,28 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
         }
 
         runBlocking {
-            writeTestToFile(request.project, flow, testContext)
+            // Check for cancellation before writing to file
+            if (indicator.isCanceled) {
+                throw ProcessCanceledException()
+            }
+
+            writeTestToFile(request.project, flow, testContext, indicator)
+
+            // Check for cancellation before verification
+            if (indicator.isCanceled) {
+                throw ProcessCanceledException()
+            }
 
             indicator.fraction = 1.0
             indicator.text = AutoDevBundle.message("intentions.chat.code.test.verify")
 
             try {
                 autoTestService.collectSyntaxError(testContext.outputFile, request.project) {
+                    // Check for cancellation before fixing syntax errors
+                    if (indicator.isCanceled) {
+                        throw ProcessCanceledException()
+                    }
+
                     autoTestService.tryFixSyntaxError(testContext.outputFile, request.project, it)
 
                     if (it.isNotEmpty()) {
@@ -146,9 +185,15 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
                         )
                         indicator.fraction = 1.0
                     } else {
-                        autoTestService.runFile(request.project, testContext.outputFile, testContext.testElement, false)
+                        // Check for cancellation before running tests
+                        if (!indicator.isCanceled) {
+                            autoTestService.runFile(request.project, testContext.outputFile, testContext.testElement, false)
+                        }
                     }
                 }
+            } catch (e: ProcessCanceledException) {
+                // Re-throw cancellation exception
+                throw e
             } catch (e: Exception) {
                 AutoDevStatusService.notifyApplication(AutoDevStatus.Ready)
                 indicator.fraction = 1.0
@@ -169,6 +214,7 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
         project: Project,
         flow: Flow<String>,
         context: TestFileContext,
+        indicator: ProgressIndicator
     ) {
         val fileEditorManager = FileEditorManager.getInstance(project)
         var editors: Array<FileEditor> = emptyArray()
@@ -183,6 +229,11 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
             val editor = fileEditorManager.selectedTextEditor
 
             flow.collect {
+                // Check for cancellation during flow collection
+                if (indicator.isCanceled) {
+                    throw ProcessCanceledException()
+                }
+
                 suggestion.append(it)
                 val codeBlocks = MarkdownCodeHelper.parseCodeFromString(suggestion.toString())
                 codeBlocks.forEach {
@@ -200,6 +251,10 @@ class TestCodeGenTask(val request: TestCodeGenRequest, displayMessage: String) :
 
         val suggestion = StringBuilder()
         flow.collect {
+            // Check for cancellation during flow collection
+            if (indicator.isCanceled) {
+                throw ProcessCanceledException()
+            }
             suggestion.append(it)
         }