fix incorrect state after indicator cancellation

JetBrains-Research · Oct 10, 2024 · f05e950 · f05e950
1 parent c1afd2b
commit f05e950
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/.../main/kotlin/org/jetbrains/research/testspark/core/generation/llm/LLMWithFeedbackCycle.kt b/.../main/kotlin/org/jetbrains/research/testspark/core/generation/llm/LLMWithFeedbackCycle.kt
@@ -147,6 +147,12 @@ class LLMWithFeedbackCycle(
                 errorMonitor,
             )
 
+            // Process stopped checking
+            if (indicator.isCanceled()) {
+                executionResult = FeedbackCycleExecutionResult.CANCELED
+                break
+            }
+
             when (response.errorCode) {
                 ResponseErrorCode.OK -> {
                     log.info { "Test suite generated successfully: ${response.testSuite!!}" }
@@ -264,6 +270,12 @@ class LLMWithFeedbackCycle(
             // saving the compilable test cases
             compilableTestCases.addAll(testCasesCompilationResult.compilableTestCases)
 
+            // Process stopped checking
+            if (indicator.isCanceled()) {
+                executionResult = FeedbackCycleExecutionResult.CANCELED
+                break
+            }
+
             if (!testCasesCompilationResult.allTestCasesCompilable && !isLastIteration(requestsCount)) {
                 log.info { "Non-compilable test suite: \n${testsPresenter.representTestSuite(generatedTestSuite!!)}" }
 

diff --git a/src/main/kotlin/org/jetbrains/research/testspark/tools/ToolUtils.kt b/src/main/kotlin/org/jetbrains/research/testspark/tools/ToolUtils.kt
@@ -109,12 +109,13 @@ object ToolUtils {
      * @return true if the process has been stopped, false otherwise
      */
     fun isProcessStopped(errorMonitor: ErrorMonitor, indicator: CustomProgressIndicator): Boolean {
-        return errorMonitor.hasErrorOccurred() || isProcessCanceled(indicator)
+        return errorMonitor.hasErrorOccurred() || isProcessCanceled(errorMonitor, indicator)
     }
 
-    fun isProcessCanceled(indicator: CustomProgressIndicator): Boolean {
+    fun isProcessCanceled(errorMonitor: ErrorMonitor, indicator: CustomProgressIndicator): Boolean {
         if (indicator.isCanceled()) {
-            // TODO: we must not stop this indicator! cancellation MAY imply stoppage
+            errorMonitor.notifyErrorOccurrence()
+            // TODO: we must not stop this indicator! cancellation MAY already imply stoppage
             //       See: https://github.com/JetBrains-Research/TestSpark/issues/375
             indicator.stop()
             return true

diff --git a/...tlin/org/jetbrains/research/testspark/tools/llm/generation/openai/OpenAIRequestManager.kt b/...tlin/org/jetbrains/research/testspark/tools/llm/generation/openai/OpenAIRequestManager.kt
@@ -58,7 +58,7 @@ class OpenAIRequestManager(project: Project) : IJRequestManager(project) {
                 // check response
                 when (val responseCode = connection.responseCode) {
                     HttpURLConnection.HTTP_OK -> {
-                        assembleLlmResponse(request, testsAssembler, indicator)
+                        assembleLlmResponse(request, testsAssembler, indicator, errorMonitor)
                     }
 
                     HttpURLConnection.HTTP_INTERNAL_ERROR -> {
@@ -115,9 +115,10 @@ class OpenAIRequestManager(project: Project) : IJRequestManager(project) {
         httpRequest: HttpRequests.Request,
         testsAssembler: TestsAssembler,
         indicator: CustomProgressIndicator,
+        errorMonitor: ErrorMonitor,
     ) {
         while (true) {
-            if (ToolUtils.isProcessCanceled(indicator)) return
+            if (ToolUtils.isProcessCanceled(errorMonitor, indicator)) return
 
             var text = httpRequest.reader.readLine()