Fix rejected files display inconsistency between VS Code and JetBrains

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

@@ -298,7 +298,6 @@ class FeatureDevController(
                 filePaths = filePaths.filter { it.zipFilePath == fileToUpdate },
                 deletedFiles = deletedFiles.filter { it.zipFilePath == fileToUpdate },
                 references = references, // Add all references (not attributed per-file)
-                messenger
             )
 
             AmazonqTelemetry.isAcceptedCodeChanges(
@@ -313,8 +312,6 @@ class FeatureDevController(
             deletedFiles.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
         }
 
-        // FIXME: This is a kludge that is hiding the fact that insertChanges is updating the file tree above this point to
-        // an incorrect state. 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:
@@ -421,14 +418,14 @@ class FeatureDevController(
                 credentialStartUrl = getStartUrl(project = context.project)
             )
 
-            // Caution: insertChanges has multiple responsibilities.
-            // The filter here results in rejected files being hidden from the tree after continuing, by design.
-            // However, it is critical that we don't hide already-accepted files. Inside insertChanges, it
-            // filters to only update the subset of passed files that aren't accepted or rejected already.
             session.insertChanges(
-                filePaths = filePaths.filter { !it.rejected },
-                deletedFiles = deletedFiles.filter { !it.rejected },
-                references = references,
+                filePaths = filePaths,
+                deletedFiles = deletedFiles,
+                references = references
+            )
+            session.updateFilesPaths(
+                filePaths = filePaths,
+                deletedFiles = deletedFiles,
                 messenger
             )
 

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

@@ -111,18 +111,29 @@ class Session(val tabID: String, val project: Project) {
         this._codeResultMessageId = messageId
     }
 
+    suspend fun updateFilesPaths(
+        filePaths: List<NewFileZipInfo>,
+        deletedFiles: List<DeletedFileInfo>,
+        messenger: MessagePublisher,
+        disableFileActions: Boolean = false,
+    ) {
+        val codeResultMessageId = this._codeResultMessageId
+        if (codeResultMessageId != null) {
+            messenger.updateFileComponent(this.tabID, filePaths, deletedFiles, codeResultMessageId, disableFileActions)
+        }
+    }
+
     /**
      * Triggered by the Insert code follow-up button to apply code changes.
      */
     suspend fun insertChanges(
         filePaths: List<NewFileZipInfo>,
         deletedFiles: List<DeletedFileInfo>,
         references: List<CodeReferenceGenerated>,
-        messenger: MessagePublisher,
     ) {
-        val selectedSourceFolder = context.selectedSourceFolder.toNioPath()
         val newFilePaths = filePaths.filter { !it.rejected && !it.changeApplied }
         val newDeletedFiles = deletedFiles.filter { !it.rejected && !it.changeApplied }
+        val selectedSourceFolder = context.selectedSourceFolder.toNioPath()
 
         runCatching {
             var insertedLines = 0
@@ -159,23 +170,39 @@ class Session(val tabID: String, val project: Project) {
             }
         }.onFailure { /* Noop on diff telemetry failure */ }
 
-        newFilePaths.forEach {
-            resolveAndCreateOrUpdateFile(selectedSourceFolder, it.zipFilePath, it.fileContent)
-            it.changeApplied = true
-        }
+        insertNewFiles(newFilePaths)
 
-        newDeletedFiles.forEach {
-            resolveAndDeleteFile(selectedSourceFolder, it.zipFilePath)
-            it.changeApplied = true
-        }
+        applyDeleteFiles(newDeletedFiles)
 
         ReferenceLogController.addReferenceLog(references, project)
 
         // Taken from https://intellij-support.jetbrains.com/hc/en-us/community/posts/206118439-Refresh-after-external-changes-to-project-structure-and-sources
         VfsUtil.markDirtyAndRefresh(true, true, true, context.selectedSourceFolder)
-        val codeResultMessageId = this._codeResultMessageId
-        if (codeResultMessageId != null) {
-            messenger.updateFileComponent(this.tabID, filePaths, deletedFiles, codeResultMessageId)
+    }
+
+// Suppressing because insertNewFiles needs to be a suspend function in order to be tested
+    @Suppress("RedundantSuspendModifier")
+    suspend fun insertNewFiles(
+        filePaths: List<NewFileZipInfo>,
+    ) {
+        val selectedSourceFolder = context.selectedSourceFolder.toNioPath()
+
+        filePaths.forEach {
+            resolveAndCreateOrUpdateFile(selectedSourceFolder, it.zipFilePath, it.fileContent)
+            it.changeApplied = true
+        }
+    }
+
+// Suppressing because applyDeleteFiles needs to be a suspend function in order to be tested
+    @Suppress("RedundantSuspendModifier")
+    suspend fun applyDeleteFiles(
+        deletedFiles: List<DeletedFileInfo>,
+    ) {
+        val selectedSourceFolder = context.selectedSourceFolder.toNioPath()
+
+        deletedFiles.forEach {
+            resolveAndDeleteFile(selectedSourceFolder, it.zipFilePath)
+            it.changeApplied = true
         }
     }
 
@@ -188,10 +215,7 @@ class Session(val tabID: String, val project: Project) {
             return
         }
 
-        val codeResultMessageId = this._codeResultMessageId
-        if (codeResultMessageId != null) {
-            messenger.updateFileComponent(this.tabID, filePaths, deletedFiles, codeResultMessageId, disableFileActions = true)
-        }
+        updateFilesPaths(filePaths, deletedFiles, messenger, disableFileActions = true)
         this._codeResultMessageId = null
     }
 

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt

@@ -241,22 +241,34 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 ),
             )
 
-            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any(), any())
+            doReturn(Unit).whenever(spySession).insertChanges(any(), any(), any())
+            doReturn(Unit).whenever(spySession).insertNewFiles(any())
+            doReturn(Unit).whenever(spySession).applyDeleteFiles(any())
 
             spySession.preloader(userMessage, messenger)
             controller.processFollowupClickedMessage(message)
 
             mockitoVerify(
                 spySession,
                 times(1),
-            ).insertChanges(listOf(newFileContents[0]), listOf(deletedFiles[0]), testReferences, messenger) // insert changes for only non rejected files
+            ).insertChanges(newFileContents, deletedFiles, testReferences) // updates for all files
             coVerifyOrder {
                 AmazonqTelemetry.isAcceptedCodeChanges(
                     amazonqNumberOfFilesAccepted = 2.0, // it should be 2 files per test setup
                     amazonqConversationId = spySession.conversationId,
                     enabled = true,
                     createTime = any(),
                 )
+
+                // insert changes for only non rejected files
+                spySession.insertNewFiles(listOf(newFileContents[0]))
+                spySession.applyDeleteFiles(listOf(deletedFiles[0]))
+
+                spySession.updateFilesPaths(
+                    filePaths = newFileContents,
+                    deletedFiles = deletedFiles,
+                    messenger
+                )
                 messenger.sendAnswer(
                     tabId = testTabId,
                     message = message("amazonqFeatureDev.code_generation.updated_code"),
@@ -468,7 +480,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 ),
             )
             doReturn(testConversationId).`when`(spySession).conversationId
-            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any(), any())
+            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any())
 
             mockkObject(AmazonqTelemetry)
             every {
@@ -491,7 +503,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
             mockitoVerify(
                 spySession,
                 times(1),
-            ).insertChanges(listOf(newFileContents[0]), listOf(), testReferences, messenger)
+            ).insertChanges(listOf(newFileContents[0]), listOf(), testReferences)
 
             // Does not continue automatically, because files are remaining:
             mockitoVerify(
@@ -525,7 +537,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 ),
             )
             doReturn(testConversationId).`when`(spySession).conversationId
-            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any(), any())
+            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any())
 
             mockkObject(AmazonqTelemetry)
             every {

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionTest.kt

@@ -83,7 +83,7 @@ class SessionTest : FeatureDevTestBase() {
         whenever(session.context.selectedSourceFolder.toNioPath()).thenReturn(Path(""))
 
         runBlocking {
-            session.insertChanges(mockNewFile, mockDeletedFile, emptyList(), messenger)
+            session.insertChanges(mockNewFile, mockDeletedFile, emptyList())
         }
 
         verify(exactly = 1) { resolveAndDeleteFile(any(), "deletedTest.ts") }