Fix rejected files display inconsistency between VS Code and JetBrains

Author: siakmun-aws

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

@@ -294,7 +294,7 @@ class FeatureDevController(
         val prevInsertAction = insertAction()
 
         if (action == "accept-change") {
-            session.insertChanges(
+            session.insertChangesWithoutUpdateFileComponents(
                 filePaths = filePaths.filter { it.zipFilePath == fileToUpdate },
                 deletedFiles = deletedFiles.filter { it.zipFilePath == fileToUpdate },
                 references = references, // Add all references (not attributed per-file)
@@ -313,8 +313,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,13 +419,9 @@ 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 },
+            session.insertChangesAndUpdateFileComponents(
+                filePaths = filePaths,
+                deletedFiles = deletedFiles,
                 references = references,
                 messenger
             )

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

@@ -114,15 +114,40 @@ class Session(val tabID: String, val project: Project) {
     /**
      * Triggered by the Insert code follow-up button to apply code changes.
      */
-    suspend fun insertChanges(
+    suspend fun insertChangesAndUpdateFileComponents(
+        filePaths: List<NewFileZipInfo>,
+        deletedFiles: List<DeletedFileInfo>,
+        references: List<CodeReferenceGenerated>,
+        messenger: MessagePublisher,
+    ) {
+        insertChangesWithoutUpdateFileComponents(filePaths, deletedFiles, references, messenger)
+        val codeResultMessageId = this._codeResultMessageId
+        if (codeResultMessageId != null) {
+            messenger.updateFileComponent(this.tabID, filePaths, deletedFiles, codeResultMessageId)
+        }
+    }
+
+    suspend fun insertChangesWithoutUpdateFileComponents(
         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 }
+        insertChanges(newFilePaths, 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)
+    }
+
+    suspend fun insertChanges(
+        filePaths: List<NewFileZipInfo>,
+        deletedFiles: List<DeletedFileInfo>,
+    ) {
+        val selectedSourceFolder = context.selectedSourceFolder.toNioPath()
 
         runCatching {
             var insertedLines = 0
@@ -159,24 +184,15 @@ class Session(val tabID: String, val project: Project) {
             }
         }.onFailure { /* Noop on diff telemetry failure */ }
 
-        newFilePaths.forEach {
+        filePaths.forEach {
             resolveAndCreateOrUpdateFile(selectedSourceFolder, it.zipFilePath, it.fileContent)
             it.changeApplied = true
         }
 
-        newDeletedFiles.forEach {
+        deletedFiles.forEach {
             resolveAndDeleteFile(selectedSourceFolder, it.zipFilePath)
             it.changeApplied = true
         }
-
-        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)
-        }
     }
 
     suspend fun disableFileList(

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

@@ -241,22 +241,24 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 ),
             )
 
-            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any(), any())
+            doReturn(Unit).`when`(spySession).insertChangesAndUpdateFileComponents(any(), any(), any(), any())
+            doReturn(Unit).`when`(spySession).insertChanges(any(), 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
+            ).insertChangesAndUpdateFileComponents(newFileContents, deletedFiles, testReferences, messenger) // 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(),
                 )
+                spySession.insertChanges(listOf(newFileContents[0]), listOf(deletedFiles[0])) // insert changes for only non-rejected files
                 messenger.sendAnswer(
                     tabId = testTabId,
                     message = message("amazonqFeatureDev.code_generation.updated_code"),
@@ -468,7 +470,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 ),
             )
             doReturn(testConversationId).`when`(spySession).conversationId
-            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any(), any())
+            doReturn(Unit).`when`(spySession).insertChangesWithoutUpdateFileComponents(any(), any(), any(), any())
 
             mockkObject(AmazonqTelemetry)
             every {
@@ -491,7 +493,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
             mockitoVerify(
                 spySession,
                 times(1),
-            ).insertChanges(listOf(newFileContents[0]), listOf(), testReferences, messenger)
+            ).insertChangesWithoutUpdateFileComponents(listOf(newFileContents[0]), listOf(), testReferences, messenger)
 
             // Does not continue automatically, because files are remaining:
             mockitoVerify(
@@ -525,7 +527,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 ),
             )
             doReturn(testConversationId).`when`(spySession).conversationId
-            doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any(), any())
+            doReturn(Unit).`when`(spySession).insertChangesWithoutUpdateFileComponents(any(), any(), any(), any())
 
             mockkObject(AmazonqTelemetry)
             every {

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

@@ -65,7 +65,7 @@ class SessionTest : FeatureDevTestBase() {
     }
 
     @Test
-    fun `test insertChanges`() {
+    fun `test insertChangesAndUpdateFileComponents`() {
         mockkStatic("com.intellij.openapi.vfs.VfsUtil")
         every { VfsUtil.markDirtyAndRefresh(true, true, true, any<VirtualFile>()) } just runs
 
@@ -83,7 +83,7 @@ class SessionTest : FeatureDevTestBase() {
         whenever(session.context.selectedSourceFolder.toNioPath()).thenReturn(Path(""))
 
         runBlocking {
-            session.insertChanges(mockNewFile, mockDeletedFile, emptyList(), messenger)
+            session.insertChangesAndUpdateFileComponents(mockNewFile, mockDeletedFile, emptyList(), messenger)
         }
 
         verify(exactly = 1) { resolveAndDeleteFile(any(), "deletedTest.ts") }