feat(amazonq): Implement file-level acceptance for /dev.

Co-authored-by: Thiago Verney <[email protected]>

Author: C Tidd

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

@@ -5,8 +5,10 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.controller
 
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.intellij.diff.DiffContentFactory
-import com.intellij.diff.DiffManager
+import com.intellij.diff.chains.SimpleDiffRequestChain
 import com.intellij.diff.contents.EmptyContent
+import com.intellij.diff.editor.ChainDiffVirtualFile
+import com.intellij.diff.editor.DiffEditorTabFilesManager
 import com.intellij.diff.requests.SimpleDiffRequest
 import com.intellij.ide.BrowserUtil
 import com.intellij.openapi.application.runInEdt
@@ -17,6 +19,7 @@ import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.wm.ToolWindowManager
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.withContext
 import software.amazon.awssdk.services.toolkittelemetry.model.Sentiment
 import software.aws.toolkits.core.utils.debug
@@ -65,7 +68,10 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Prepar
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.selectFolder
+import software.aws.toolkits.jetbrains.services.codewhisperer.util.content
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.FeedbackComment
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
 import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
@@ -86,13 +92,19 @@ class FeatureDevController(
     val messenger = context.messagesFromAppToUi
     val toolWindow = ToolWindowManager.getInstance(context.project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID)
 
+    private val diffVirtualFiles = mutableMapOf<String, ChainDiffVirtualFile>()
+
     override suspend fun processPromptChatMessage(message: IncomingFeatureDevMessage.ChatPrompt) {
         handleChat(
             tabId = message.tabId,
             message = message.chatMessage
         )
     }
 
+    override suspend fun processStoreCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) {
+        storeCodeResultMessageId(message)
+    }
+
     override suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse) {
         handleStopMessage(message)
     }
@@ -126,6 +138,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) {
@@ -192,6 +205,18 @@ class FeatureDevController(
         }
     }
 
+    private fun putDiff(filePath: String, request: SimpleDiffRequest) {
+        // Close any existing diff and open a new diff, as the diff virtual file does not appear to allow replacing content directly:
+        val existingDiff = diffVirtualFiles[filePath]
+        if (existingDiff != null) {
+            FileEditorManager.getInstance(context.project).closeFile(existingDiff)
+        }
+
+        val newDiff = ChainDiffVirtualFile(SimpleDiffRequestChain(request), filePath)
+        DiffEditorTabFilesManager.getInstance(context.project).showDiffFile(newDiff, true)
+        diffVirtualFiles[filePath] = newDiff
+    }
+
     override suspend fun processOpenDiff(message: IncomingFeatureDevMessage.OpenDiff) {
         val session = getSessionInfo(message.tabId)
 
@@ -223,9 +248,7 @@ class FeatureDevController(
                         DiffContentFactory.getInstance().create(newFileContent)
                     }
 
-                    val request = SimpleDiffRequest(message.filePath, leftDiffContent, rightDiffContent, null, null)
-
-                    DiffManager.getInstance().showDiff(project, request)
+                    putDiff(message.filePath, SimpleDiffRequest(message.filePath, leftDiffContent, rightDiffContent, null, null))
                 }
             }
             else -> {
@@ -244,6 +267,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()
@@ -254,11 +278,69 @@ class FeatureDevController(
             }
         }
 
-        // 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 }
+        fun insertAction(): InsertAction =
+            if (filePaths.all { it.changeApplied } && deletedFiles.all { it.changeApplied }) {
+                InsertAction.AUTO_CONTINUE
+            } else if (filePaths.all { it.changeApplied || it.rejected } && deletedFiles.all { it.changeApplied || it.rejected }) {
+                InsertAction.CONTINUE
+            } else if (filePaths.any { it.changeApplied || it.rejected } || deletedFiles.any { it.changeApplied || it.rejected }) {
+                InsertAction.REMAINING
+            } else {
+                InsertAction.ALL
+            }
+
+        val prevInsertAction = insertAction()
+
+        if (action == "accept-change") {
+            session.insertChanges(
+                filePaths = filePaths.filter { it.zipFilePath == fileToUpdate },
+                deletedFiles = deletedFiles.filter { it.zipFilePath == fileToUpdate },
+                references = emptyList(),
+                messenger
+            )
+
+            AmazonqTelemetry.isAcceptedCodeChanges(
+                amazonqNumberOfFilesAccepted = 1.0,
+                amazonqConversationId = session.conversationId,
+                enabled = true,
+                credentialStartUrl = getStartUrl(project = context.project)
+            )
+        } 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 }
+        }
 
+        // Update the state of the tree view:
         messenger.updateFileComponent(message.tabId, filePaths, deletedFiles, messageId)
+
+        // Then, if the accepted file is not a deletion, open a diff to show the changes are applied:
+        if (action == "accept-change" && deletedFiles.none { it.zipFilePath == fileToUpdate }) {
+            var pollAttempt = 0
+            val pollDelayMs = 10L
+            while (pollAttempt < 5) {
+                val file = VfsUtil.findRelativeFile(message.filePath, session.context.selectedSourceFolder)
+                // Wait for the file to be created and/or updated to the new content:
+                if (file != null && file.content() == filePaths.find { it.zipFilePath == fileToUpdate }?.fileContent) {
+                    // Open a diff, showing the changes have been applied and the file now has identical left/right state:
+                    this.processOpenDiff(IncomingFeatureDevMessage.OpenDiff(message.tabId, fileToUpdate, false))
+                    break
+                } else {
+                    pollAttempt++
+                    delay(pollDelayMs)
+                }
+            }
+        }
+
+        val nextInsertAction = insertAction()
+
+        if (nextInsertAction == InsertAction.AUTO_CONTINUE) {
+            // Insert remaining changes (noop, as there are none), and advance to the next prompt:
+            insertCode(message.tabId)
+        } else if (nextInsertAction != prevInsertAction) {
+            // Update the action displayed to the customer based on the current state:
+            messenger.sendSystemPrompt(message.tabId, getFollowUpOptions(session.sessionState.phase, nextInsertAction))
+        }
     }
 
     private suspend fun newTabOpened(tabId: String) {
@@ -325,8 +407,12 @@ class FeatureDevController(
                 }
             }
 
+            val rejectedFilesCount = filePaths.count { it.rejected } + deletedFiles.count { it.rejected }
+            val acceptedFilesCount = filePaths.count { it.changeApplied } + filePaths.count { it.changeApplied }
+            val remainingFilesCount = filePaths.count() + deletedFiles.count() - acceptedFilesCount - rejectedFilesCount
+
             AmazonqTelemetry.isAcceptedCodeChanges(
-                amazonqNumberOfFilesAccepted = (filePaths.filterNot { it.rejected }.size + deletedFiles.filterNot { it.rejected }.size) * 1.0,
+                amazonqNumberOfFilesAccepted = remainingFilesCount * 1.0,
                 amazonqConversationId = session.conversationId,
                 enabled = true,
                 credentialStartUrl = getStartUrl(project = context.project)
@@ -335,7 +421,8 @@ class FeatureDevController(
             session.insertChanges(
                 filePaths = filePaths.filterNot { it.rejected },
                 deletedFiles = deletedFiles.filterNot { it.rejected },
-                references = references
+                references = references,
+                messenger
             )
 
             messenger.sendAnswer(
@@ -377,8 +464,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 +489,7 @@ class FeatureDevController(
     }
 
     private suspend fun closeSession(tabId: String) {
+        this.disablePreviousFileList(tabId)
         messenger.sendAnswer(
             tabId = tabId,
             messageType = FeatureDevMessageType.Answer,
@@ -503,7 +594,7 @@ class FeatureDevController(
                     tabId = tabId,
                     followUp = listOf(
                         FollowUp(
-                            pillText = message("amazonqFeatureDev.follow_up.insert_code"),
+                            pillText = message("amazonqFeatureDev.follow_up.insert_all_code"),
                             type = FollowUpTypes.INSERT_CODE,
                             icon = FollowUpIcons.Ok,
                             status = FollowUpStatusType.Success,
@@ -546,11 +637,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

@@ -23,6 +23,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Prepar
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionState
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.resources.message
@@ -88,7 +89,7 @@ suspend fun FeatureDevController.onCodeGeneration(
         }
 
         // Atm this is the only possible path as codegen is mocked to return empty.
-        if (filePaths.size or deletedFiles.size == 0) {
+        if (filePaths.isEmpty() && deletedFiles.isEmpty()) {
             messenger.sendAnswer(
                 tabId = tabId,
                 messageType = FeatureDevMessageType.Answer,
@@ -132,7 +133,7 @@ suspend fun FeatureDevController.onCodeGeneration(
             )
         }
 
-        messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase))
+        messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase, InsertAction.ALL))
 
         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"