fix(/dev): source folder modification now allows any sub-folder (#4489)

sannicm · web-flow · commit 6759455963e3 · 2024-06-12T16:39:36.000+02:00
The source folder modification happens when the customers repository is too big for uploading and /dev allows the customer to select a new source folder from within their working project. Until now, this check was faulty as it would only check for one level of sub-folder from the root, now it is any level. And of course, folders outside the workspace are not allowed. Workspace being for VSCode the project's root.
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevConstants.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevConstants.kt
@@ -16,3 +16,13 @@ const val DEFAULT_RETRY_LIMIT = 0
 
 // Max allowed size for a repository in bytes
 const val MAX_PROJECT_SIZE_BYTES: Long = 200 * 1024 * 1024
+
+enum class ModifySourceFolderErrorReason(
+    private val reasonText: String
+) {
+    ClosedBeforeSelection("ClosedBeforeSelection"),
+    NotInWorkspaceFolder("NotInWorkspaceFolder"),
+    ;
+
+    override fun toString(): String = reasonText
+}
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt
@@ -13,8 +13,6 @@ import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.Caret
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.fileChooser.FileChooser
-import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.wm.ToolWindowManager
@@ -33,6 +31,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeIterationL
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.DEFAULT_RETRY_LIMIT
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.InboundAppMessagesHandler
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ModifySourceFolderErrorReason
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MonthlyConversationLimitError
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.PlanIterationLimitError
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.createUserFacingErrorMessage
@@ -61,6 +60,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Sessio
 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.getFollowUpOptions
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.selectFolder
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
 import software.aws.toolkits.jetbrains.ui.feedback.FeatureDevFeedbackDialog
 import software.aws.toolkits.resources.message
@@ -192,7 +192,7 @@ class FeatureDevController(
         when (sessionState) {
             is PrepareCodeGenerationState -> {
                 runInEdt {
-                    val existingFile = VfsUtil.findRelativeFile(message.filePath, session.context.projectRoot)
+                    val existingFile = VfsUtil.findRelativeFile(message.filePath, session.context.selectedSourceFolder)
 
                     val leftDiffContent = if (existingFile == null) {
                         EmptyContent()
@@ -602,8 +602,8 @@ class FeatureDevController(
 
     private suspend fun modifyDefaultSourceFolder(tabId: String) {
         val session = getSessionInfo(tabId)
-        val uri = session.context.projectRoot
-        val fileChooserDescriptor = FileChooserDescriptorFactory.createSingleFolderDescriptor()
+        val currentSourceFolder = session.context.selectedSourceFolder
+        val projectRoot = session.context.projectRoot
 
         val modifyFolderFollowUp = FollowUp(
             pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"),
@@ -612,11 +612,10 @@ class FeatureDevController(
         )
 
         var result: Result = Result.Failed
-        var reason: String? = null
+        var reason: ModifySourceFolderErrorReason? = null
 
         withContext(EDT) {
-            val selectedFolder = FileChooser.chooseFile(fileChooserDescriptor, context.project, uri)
-
+            val selectedFolder = selectFolder(context.project, currentSourceFolder)
             // No folder was selected
             if (selectedFolder == null) {
                 logger.info { "Cancelled dialog and not selected any folder" }
@@ -626,12 +625,12 @@ class FeatureDevController(
                     followUp = listOf(modifyFolderFollowUp),
                 )
 
-                reason = "ClosedBeforeSelection"
+                reason = ModifySourceFolderErrorReason.ClosedBeforeSelection
                 return@withContext
             }
 
             // The folder is not in the workspace
-            if (selectedFolder.parent.path != uri.path) {
+            if (!selectedFolder.path.startsWith(projectRoot.path)) {
                 logger.info { "Selected folder not in workspace: ${selectedFolder.path}" }
 
                 messenger.sendAnswer(
@@ -645,13 +644,13 @@ class FeatureDevController(
                     followUp = listOf(modifyFolderFollowUp),
                 )
 
-                reason = "NotInWorkspaceFolder"
+                reason = ModifySourceFolderErrorReason.NotInWorkspaceFolder
                 return@withContext
             }
 
             logger.info { "Selected correct folder inside workspace: ${selectedFolder.path}" }
 
-            session.context.projectRoot = selectedFolder
+            session.context.selectedSourceFolder = selectedFolder
             result = Result.Succeeded
 
             messenger.sendAnswer(
@@ -665,7 +664,7 @@ class FeatureDevController(
             amazonqConversationId = session.conversationId,
             credentialStartUrl = getStartUrl(project = context.project),
             result = result,
-            reason = reason
+            reason = reason?.toString()
         )
     }
 
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt
@@ -97,16 +97,16 @@ class Session(val tabID: String, val project: Project) {
      * Triggered by the Insert code follow-up button to apply code changes.
      */
     fun insertChanges(filePaths: List<NewFileZipInfo>, deletedFiles: List<DeletedFileInfo>, references: List<CodeReferenceGenerated>) {
-        val projectRootPath = context.projectRoot.toNioPath()
+        val selectedSourceFolder = context.selectedSourceFolder.toNioPath()
 
-        filePaths.forEach { resolveAndCreateOrUpdateFile(projectRootPath, it.zipFilePath, it.fileContent) }
+        filePaths.forEach { resolveAndCreateOrUpdateFile(selectedSourceFolder, it.zipFilePath, it.fileContent) }
 
-        deletedFiles.forEach { resolveAndDeleteFile(projectRootPath, it.zipFilePath) }
+        deletedFiles.forEach { resolveAndDeleteFile(selectedSourceFolder, it.zipFilePath) }
 
         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.projectRoot)
+        VfsUtil.markDirtyAndRefresh(true, true, true, context.selectedSourceFolder)
     }
 
     suspend fun send(msg: String): Interaction {
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FileUtils.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FileUtils.kt
@@ -3,6 +3,10 @@
 
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util
 
+import com.intellij.openapi.fileChooser.FileChooser
+import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import java.nio.file.Path
 import kotlin.io.path.createDirectories
 import kotlin.io.path.deleteIfExists
@@ -24,3 +28,8 @@ fun resolveAndDeleteFile(projectRootPath: Path, relativePath: String) {
     val filePath = projectRootPath.resolve(relativePath)
     filePath.deleteIfExists()
 }
+
+fun selectFolder(project: Project, openOn: VirtualFile): VirtualFile? {
+    val fileChooserDescriptor = FileChooserDescriptorFactory.createSingleFolderDescriptor()
+    return FileChooser.chooseFile(fileChooserDescriptor, project, openOn)
+}
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt
@@ -3,6 +3,7 @@
 
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.controller
 
+import com.intellij.testFramework.LightVirtualFile
 import com.intellij.testFramework.RuleChain
 import com.intellij.testFramework.replaceService
 import io.mockk.coEvery
@@ -56,6 +57,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Sessio
 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.getFollowUpOptions
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.selectFolder
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.AmazonqTelemetry
@@ -394,4 +396,89 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         newFileContentsCopy[0].rejected = !newFileContentsCopy[0].rejected
         coVerify { messenger.updateFileComponent(testTabId, newFileContentsCopy, deletedFiles, "") }
     }
+
+    @Test
+    fun `test modifyDefaultSourceFolder customer does not select a folder`() = runTest {
+        val followUp = FollowUp(FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, pillText = "Modify default source folder")
+        val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command")
+
+        whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse)
+        whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
+
+        mockkStatic("software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FileUtilsKt")
+        every { selectFolder(any(), any()) } returns null
+
+        spySession.preloader(userMessage, messenger)
+        controller.processFollowupClickedMessage(message)
+
+        coVerifyOrder {
+            messenger.sendSystemPrompt(
+                tabId = testTabId,
+                followUp = listOf(
+                    FollowUp(
+                        pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"),
+                        type = FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER,
+                        status = FollowUpStatusType.Info,
+                    )
+                )
+            )
+        }
+    }
+
+    @Test
+    fun `test modifyDefaultSourceFolder customer selects a folder outside the workspace`() = runTest {
+        val followUp = FollowUp(FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, pillText = "Modify default source folder")
+        val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command")
+
+        whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse)
+        whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
+
+        mockkStatic("software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FileUtilsKt")
+        every { selectFolder(any(), any()) } returns LightVirtualFile("/path")
+
+        spySession.preloader(userMessage, messenger)
+        controller.processFollowupClickedMessage(message)
+
+        coVerifyOrder {
+            messenger.sendAnswer(
+                tabId = testTabId,
+                messageType = FeatureDevMessageType.Answer,
+                message = message("amazonqFeatureDev.follow_up.incorrect_source_folder")
+            )
+            messenger.sendSystemPrompt(
+                tabId = testTabId,
+                followUp = listOf(
+                    FollowUp(
+                        pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"),
+                        type = FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER,
+                        status = FollowUpStatusType.Info,
+                    )
+                )
+            )
+        }
+    }
+
+    @Test
+    fun `test modifyDefaultSourceFolder customer selects a correct sub folder`() = runTest {
+        val followUp = FollowUp(FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, pillText = "Modify default source folder")
+        val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command")
+
+        whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse)
+        whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
+
+        val folder = LightVirtualFile("${spySession.context.projectRoot.name}/path/to/sub/folder")
+        mockkStatic("software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FileUtilsKt")
+        every { selectFolder(any(), any()) } returns folder
+
+        spySession.preloader(userMessage, messenger)
+        controller.processFollowupClickedMessage(message)
+
+        coVerify {
+            messenger.sendAnswer(
+                tabId = testTabId,
+                messageType = FeatureDevMessageType.Answer,
+                message = message("amazonqFeatureDev.follow_up.modified_source_folder", folder.path)
+            )
+        }
+    }
 }
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionTest.kt
@@ -91,8 +91,8 @@ class SessionTest : FeatureDevTestBase() {
         val mockNewFile = listOf(NewFileZipInfo("test.ts", "testContent", false))
         val mockDeletedFile = listOf(DeletedFileInfo("deletedTest.ts", false))
 
-        session.context.projectRoot = mock()
-        whenever(session.context.projectRoot.toNioPath()).thenReturn(Path(""))
+        session.context.selectedSourceFolder = mock()
+        whenever(session.context.selectedSourceFolder.toNioPath()).thenReturn(Path(""))
 
         session.insertChanges(mockNewFile, mockDeletedFile, emptyList())
 
diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/FeatureDevSessionContext.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/FeatureDevSessionContext.kt
@@ -66,9 +66,13 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
         "dist/"
     ).map { Regex(it) }
 
-    private var _projectRoot = project.guessProjectDir() ?: error("Cannot guess base directory for project ${project.name}")
+    // projectRoot: is the directory where the project is located when selected to open a project.
+    val projectRoot = project.guessProjectDir() ?: error("Cannot guess base directory for project ${project.name}")
+
+    // selectedSourceFolder": is the directory selected in replacement of the root, this happens when the project is too big to bundle for uploading.
+    private var _selectedSourceFolder = projectRoot
     private var ignorePatternsWithGitIgnore = emptyList<Regex>()
-    private val gitIgnoreFile = File(projectRoot.path, ".gitignore")
+    private val gitIgnoreFile = File(selectedSourceFolder.path, ".gitignore")
 
     init {
         ignorePatternsWithGitIgnore = (ignorePatterns + parseGitIgnore().map { Regex(it) }).toList()
@@ -77,7 +81,7 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
     fun getProjectZip(): ZipCreationResult {
         val zippedProject = runBlocking {
             withBackgroundProgress(project, message("amazonqFeatureDev.create_plan.background_progress_title")) {
-                zipFiles(projectRoot)
+                zipFiles(selectedSourceFolder)
             }
         }
         val checkSum256: String = Base64.getEncoder().encodeToString(DigestUtils.sha256(FileInputStream(zippedProject)))
@@ -162,11 +166,11 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
         .replace("*", ".*")
         .let { if (it.endsWith("/")) "$it?" else it } // Handle directory-specific patterns by optionally matching trailing slash
 
-    var projectRoot: VirtualFile
+    var selectedSourceFolder: VirtualFile
         set(newRoot) {
-            _projectRoot = newRoot
+            _selectedSourceFolder = newRoot
         }
-        get() = _projectRoot
+        get() = _selectedSourceFolder
 }
 
 data class ZipCreationResult(val payload: File, val checksum: String, val contentLength: Long)
