fix(amazonq): Fix code generation error when cancelling /doc task

Author: aggagen

.changes/next-release/bugfix-b99698f8-eb82-4820-b5f7-002e0ed82827.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "/doc: Fix code generation error when cancelling a documentation task"
+}

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocController.kt

@@ -239,6 +239,9 @@ class DocController(
         if (session.sessionState.token?.token !== null) {
             session.sessionState.token?.cancel()
         }
+
+        docGenerationTask.reset()
+        newTask(message.tabId)
     }
 
     private suspend fun updateDocumentation(tabId: String) {
@@ -588,6 +591,8 @@ class DocController(
             message = message("amazonqFeatureDev.chat_message.ask_for_new_task")
         )
 
+        messenger.sendUpdatePromptProgress(tabId, null)
+
         messenger.sendUpdatePlaceholder(
             tabId = tabId,
             newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation")
@@ -824,9 +829,11 @@ class DocController(
                 else -> emptyList()
             }
 
-            processOpenDiff(
-                message = IncomingDocMessage.OpenDiff(tabId = tabId, filePath = filePaths[0].zipFilePath, deleted = false)
-            )
+            if (!filePaths.isEmpty()) {
+                processOpenDiff(
+                    message = IncomingDocMessage.OpenDiff(tabId = tabId, filePath = filePaths[0].zipFilePath, deleted = false)
+                )
+            }
         } catch (err: Exception) {
             processErrorChatMessage(err, session, tabId)
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocControllerExtensions.kt

@@ -44,6 +44,13 @@ suspend fun DocController.onCodeGeneration(session: DocSession, message: String,
 
         session.send(sessionMessage) // Trigger code generation
 
+        if (session.sessionState.token
+                ?.token
+                ?.isCancellationRequested() == true
+        ) {
+            return
+        }
+
         val state = session.sessionState
 
         var filePaths: List<NewFileZipInfo> = emptyList()

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocGenerationState.kt

@@ -17,10 +17,6 @@ import software.aws.toolkits.jetbrains.services.amazonqDoc.controller.Mode
 import software.aws.toolkits.jetbrains.services.amazonqDoc.controller.docGenerationProgressMessage
 import software.aws.toolkits.jetbrains.services.amazonqDoc.docServiceError
 import software.aws.toolkits.jetbrains.services.amazonqDoc.inProgress
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.DocMessageType
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUp
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUpTypes
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.sendAnswer
 import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.sendAnswerPart
 import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.sendUpdatePromptProgress
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeGenerationResult
@@ -144,24 +140,6 @@ private suspend fun DocGenerationState.generateCode(codeGenerationId: String, to
 
     repeat(pollCount) {
         if (token?.token?.isCancellationRequested() == true) {
-            // This should be switched to newTask or something. Looks different than previously and may need to clean up previous run
-            messenger.sendUpdatePromptProgress(tabId = tabID, null)
-            messenger.sendAnswer(
-                messageType = DocMessageType.SystemPrompt,
-                tabId = tabID,
-                followUp = listOf(
-                    FollowUp(
-                        pillText = message("amazonqDoc.prompt.create"),
-                        prompt = message("amazonqDoc.prompt.create"),
-                        type = FollowUpTypes.CREATE_DOCUMENTATION,
-                    ),
-                    FollowUp(
-                        pillText = message("amazonqDoc.prompt.update"),
-                        prompt = message("amazonqDoc.prompt.update"),
-                        type = FollowUpTypes.UPDATE_DOCUMENTATION,
-                    )
-                )
-            )
             return CodeGenerationResult(emptyList(), emptyList(), emptyList())
         }