Draft: File level acceptance

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt

@@ -42,7 +42,8 @@ class FeatureDevApp : AmazonQApp {
             "insert_code_at_cursor_position" to IncomingFeatureDevMessage.InsertCodeAtCursorPosition::class,
             "open-diff" to IncomingFeatureDevMessage.OpenDiff::class,
             "file-click" to IncomingFeatureDevMessage.FileClicked::class,
-            "stop-response" to IncomingFeatureDevMessage.StopResponse::class
+            "stop-response" to IncomingFeatureDevMessage.StopResponse::class,
+            "store-code-result-message-id" to IncomingFeatureDevMessage.StoreMessageIdMessage::class
         )
 
         scope.launch {
@@ -84,6 +85,7 @@ class FeatureDevApp : AmazonQApp {
             is IncomingFeatureDevMessage.OpenDiff -> inboundAppMessagesHandler.processOpenDiff(message)
             is IncomingFeatureDevMessage.FileClicked -> inboundAppMessagesHandler.processFileClicked(message)
             is IncomingFeatureDevMessage.StopResponse -> inboundAppMessagesHandler.processStopMessage(message)
+            is IncomingFeatureDevMessage.StoreMessageIdMessage -> inboundAppMessagesHandler.processStoreCodeResultMessageId(message)
         }
     }
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/InboundAppMessagesHandler.kt

@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.IncomingFeatureDevMessage
 
 interface InboundAppMessagesHandler {
+
     suspend fun processPromptChatMessage(message: IncomingFeatureDevMessage.ChatPrompt)
     suspend fun processNewTabCreatedMessage(message: IncomingFeatureDevMessage.NewTabCreated)
     suspend fun processTabRemovedMessage(message: IncomingFeatureDevMessage.TabRemoved)
@@ -18,4 +19,5 @@ interface InboundAppMessagesHandler {
     suspend fun processOpenDiff(message: IncomingFeatureDevMessage.OpenDiff)
     suspend fun processFileClicked(message: IncomingFeatureDevMessage.FileClicked)
     suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse)
+    suspend fun processStoreCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage)
 }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt

@@ -93,6 +93,10 @@ class FeatureDevController(
         )
     }
 
+    override suspend fun processStoreCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) {
+        storeCodeResultMessageId(message)
+    }
+
     override suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse) {
         handleStopMessage(message)
     }
@@ -126,6 +130,7 @@ class FeatureDevController(
 
     override suspend fun processChatItemVotedMessage(message: IncomingFeatureDevMessage.ChatItemVotedMessage) {
         logger.debug { "$FEATURE_NAME: Processing ChatItemVotedMessage: $message" }
+        this.disablePreviousFileList(message.tabId)
 
         val session = chatSessionStorage.getSession(message.tabId, context.project)
         when (message.vote) {
@@ -244,6 +249,7 @@ class FeatureDevController(
         val fileToUpdate = message.filePath
         val session = getSessionInfo(message.tabId)
         val messageId = message.messageId
+        val action = message.actionName
 
         var filePaths: List<NewFileZipInfo> = emptyList()
         var deletedFiles: List<DeletedFileInfo> = emptyList()
@@ -253,12 +259,22 @@ class FeatureDevController(
                 deletedFiles = state.deletedFiles
             }
         }
-
-        // Mark the file as rejected or not depending on the previous state
-        filePaths.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
-        deletedFiles.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
+        if (action == "accept-change") {
+            session.insertChanges(
+                filePaths = filePaths.filter { it.zipFilePath == fileToUpdate },
+                deletedFiles = deletedFiles.filter { it.zipFilePath == fileToUpdate },
+                references = emptyList(),
+                messenger
+            )
+        } else {
+            // Mark the file as rejected or not depending on the previous state
+            filePaths.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
+            deletedFiles.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
+        }
 
         messenger.updateFileComponent(message.tabId, filePaths, deletedFiles, messageId)
+
+        //TODO: session.acceptCodeMessageId
     }
 
     private suspend fun newTabOpened(tabId: String) {
@@ -335,9 +351,13 @@ class FeatureDevController(
             session.insertChanges(
                 filePaths = filePaths.filterNot { it.rejected },
                 deletedFiles = deletedFiles.filterNot { it.rejected },
-                references = references
+                references = references,
+                messenger
             )
 
+            //TODO: if (session.acceptCodeMessageId) {
+
+
             messenger.sendAnswer(
                 tabId = tabId,
                 message = message("amazonqFeatureDev.code_generation.updated_code"),
@@ -377,8 +397,11 @@ class FeatureDevController(
     }
 
     private suspend fun newTask(tabId: String, isException: Boolean? = false) {
+        this.disablePreviousFileList(tabId)
+
         val session = getSessionInfo(tabId)
         val sessionLatency = System.currentTimeMillis() - session.sessionStartTime
+
         AmazonqTelemetry.endChat(
             amazonqConversationId = session.conversationId,
             amazonqEndOfTheConversationLatency = sessionLatency.toDouble(),
@@ -399,6 +422,7 @@ class FeatureDevController(
     }
 
     private suspend fun closeSession(tabId: String) {
+        this.disablePreviousFileList(tabId)
         messenger.sendAnswer(
             tabId = tabId,
             messageType = FeatureDevMessageType.Answer,
@@ -546,11 +570,28 @@ class FeatureDevController(
         }
     }
 
+    private suspend fun disablePreviousFileList(tabId: String) {
+        val session = getSessionInfo(tabId)
+        when (val sessionState = session.sessionState) {
+            is PrepareCodeGenerationState -> {
+                session.disableFileList(sessionState.filePaths, sessionState.deletedFiles, messenger)
+            }
+        }
+    }
+
+    private fun storeCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) {
+        val tabId = message.tabId
+        val session = getSessionInfo(tabId)
+        session.storeCodeResultMessageId(message)
+    }
+
     private suspend fun handleChat(
         tabId: String,
         message: String,
     ) {
         var session: Session? = null
+
+        this.disablePreviousFileList(tabId)
         try {
             logger.debug { "$FEATURE_NAME: Processing message: $message" }
             session = getSessionInfo(tabId)

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerExtensions.kt

@@ -16,6 +16,8 @@
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendCodeResult
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendSystemPrompt
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateChatAnswer
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateFileComponent
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo
@@ -88,7 +90,7 @@
         }
 
         // Atm this is the only possible path as codegen is mocked to return empty.
-        if (filePaths.size or deletedFiles.size == 0) {
+        if (filePaths.size == 0 && deletedFiles.size == 0) {
             messenger.sendAnswer(
                 tabId = tabId,
                 messageType = FeatureDevMessageType.Answer,
@@ -132,7 +134,11 @@
             )
         }
 
-        messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase))
+        if (filePaths.any { it.rejected or it.changeApplied } or deletedFiles.any { it.rejected or it.changeApplied }) {
+            messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase, "Accept remaining changes"))
+        } else {
+            messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase, null))
+        }
 
         messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation"))
     } finally {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessage.kt

@@ -24,6 +24,12 @@ sealed interface IncomingFeatureDevMessage : FeatureDevBaseMessage {
         @JsonProperty("tabID") val tabId: String,
     ) : IncomingFeatureDevMessage
 
+    data class StoreMessageIdMessage(
+        @JsonProperty("tabID") val tabId: String,
+        val command: String,
+        val messageId: String?,
+    ) : IncomingFeatureDevMessage
+
     data class NewTabCreated(
         val command: String,
         @JsonProperty("tabID") val tabId: String,
@@ -149,6 +155,7 @@ data class FileComponent(
     val filePaths: List<NewFileZipInfo>,
     val deletedFiles: List<DeletedFileInfo>,
     val messageId: String,
+    val disableFileActions: Boolean = false,
 ) : UiMessage(
     tabId = tabId,
     type = "updateFileComponent"
@@ -198,6 +205,7 @@ data class CodeResultMessage(
     val filePaths: List<NewFileZipInfo>,
     val deletedFiles: List<DeletedFileInfo>,
     val references: List<CodeReference>,
+    val messageId: String?
 ) : UiMessage(
     tabId = tabId,
     type = "codeResultMessage"