fix(amazonq): /doc add suppoort for uploading architecture diagrams (#…

…5357) Co-authored-by: Viktor Shesternyak <[email protected]>
aws · Feb 18, 2025 · 93b3b5a · 93b3b5a
1 parent e9c40fa
commit 93b3b5a
Show file tree

Hide file tree

Showing 8 changed files with 127 additions and 72 deletions.
diff --git a/.changes/next-release/feature-bfb6687a-a302-4968-8e7e-8fa0d4ef6c53.json b/.changes/next-release/feature-bfb6687a-a302-4968-8e7e-8fa0d4ef6c53.json
@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Amazon Q /doc: support making changes to architecture diagrams"
+}
diff --git a/...tbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocConstants.kt b/...tbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocConstants.kt
@@ -3,11 +3,6 @@
 
 package software.aws.toolkits.jetbrains.services.amazonqDoc
 
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUp
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUpStatusType
-import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUpTypes
-import software.aws.toolkits.resources.message
-
 const val FEATURE_EVALUATION_PRODUCT_NAME = "DocGeneration"
 
 const val FEATURE_NAME = "Amazon Q Documentation Generation"
@@ -21,25 +16,8 @@ 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
-}
-
-val NEW_SESSION_FOLLOWUPS: List<FollowUp> = listOf(
-    FollowUp(
-        pillText = message("amazonqDoc.prompt.reject.new_task"),
-        type = FollowUpTypes.NEW_TASK,
-        status = FollowUpStatusType.Info
-    ),
-    FollowUp(
-        pillText = message("amazonqDoc.prompt.reject.close_session"),
-        type = FollowUpTypes.CLOSE_SESSION,
-        status = FollowUpStatusType.Info
-    )
-)
+const val INFRA_DIAGRAM_PREFIX = "infra."
+const val DIAGRAM_SVG_EXT = "svg"
+const val DIAGRAM_DOT_EXT = "dot"
+val SUPPORTED_DIAGRAM_EXT_SET: Set<String> = setOf(DIAGRAM_SVG_EXT, DIAGRAM_DOT_EXT)
+val SUPPORTED_DIAGRAM_FILE_NAME_SET: Set<String> = SUPPORTED_DIAGRAM_EXT_SET.map { INFRA_DIAGRAM_PREFIX + it }.toSet()
diff --git a/...unity/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocController.kt b/...unity/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocController.kt
@@ -12,12 +12,15 @@ import com.intellij.diff.requests.SimpleDiffRequest
 import com.intellij.diff.util.DiffUserDataKeys
 import com.intellij.ide.BrowserUtil
 import com.intellij.openapi.application.runInEdt
+import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.wm.ToolWindowManager
+import com.intellij.testFramework.LightVirtualFile
 import kotlinx.coroutines.withContext
+import org.intellij.images.fileTypes.impl.SvgFileType
 import software.amazon.awssdk.services.codewhispererruntime.model.DocFolderLevel
 import software.amazon.awssdk.services.codewhispererruntime.model.DocInteractionType
 import software.amazon.awssdk.services.codewhispererruntime.model.DocUserDecision
@@ -33,6 +36,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitConte
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory
 import software.aws.toolkits.jetbrains.services.amazonqDoc.DEFAULT_RETRY_LIMIT
+import software.aws.toolkits.jetbrains.services.amazonqDoc.DIAGRAM_SVG_EXT
 import software.aws.toolkits.jetbrains.services.amazonqDoc.DocException
 import software.aws.toolkits.jetbrains.services.amazonqDoc.FEATURE_NAME
 import software.aws.toolkits.jetbrains.services.amazonqDoc.InboundAppMessagesHandler
@@ -374,45 +378,51 @@ class DocController(
 
     override suspend fun processOpenDiff(message: IncomingDocMessage.OpenDiff) {
         val session = getSessionInfo(message.tabId)
-
-        val project = context.project
         val sessionState = session.sessionState
 
-        when (sessionState) {
-            is PrepareDocGenerationState -> {
-                runInEdt {
-                    val existingFile = VfsUtil.findRelativeFile(message.filePath, session.context.selectedSourceFolder)
-
-                    val leftDiffContent = if (existingFile == null) {
-                        EmptyContent()
-                    } else {
-                        DiffContentFactory.getInstance().create(project, existingFile)
-                    }
-
-                    val newFileContent = sessionState.filePaths.find { it.zipFilePath == message.filePath }?.fileContent
+        if (sessionState !is PrepareDocGenerationState) {
+            logger.error { "$FEATURE_NAME: OpenDiff event is received for a conversation that has ${session.sessionState.phase} phase" }
+            messenger.sendError(
+                tabId = message.tabId,
+                errMessage = message("amazonqFeatureDev.exception.open_diff_failed"),
+                retries = 0,
+                conversationId = session.conversationIdUnsafe
+            )
+            return
+        }
 
-                    val rightDiffContent = if (message.deleted || newFileContent == null) {
-                        EmptyContent()
-                    } else {
-                        DiffContentFactory.getInstance().create(newFileContent)
-                    }
+        runInEdt {
+            val newFileContent = sessionState.filePaths.find { it.zipFilePath == message.filePath }?.fileContent
 
-                    val request = SimpleDiffRequest(message.filePath, leftDiffContent, rightDiffContent, null, null)
-                    request.putUserData(DiffUserDataKeys.FORCE_READ_ONLY, true)
+            val isSvgFile = message.filePath.lowercase().endsWith(".".plus(DIAGRAM_SVG_EXT))
+            if (isSvgFile && newFileContent != null) {
+                // instead of diff display generated svg in edit/preview window
+                val inMemoryFile = LightVirtualFile(
+                    message.filePath,
+                    SvgFileType.INSTANCE,
+                    newFileContent
+                )
+                inMemoryFile.isWritable = false
+                FileEditorManager.getInstance(context.project).openFile(inMemoryFile, true)
+            } else {
+                val existingFile = VfsUtil.findRelativeFile(message.filePath, session.context.selectedSourceFolder)
+                val leftDiffContent = if (existingFile == null) {
+                    EmptyContent()
+                } else {
+                    DiffContentFactory.getInstance().create(context.project, existingFile)
+                }
 
-                    val newDiff = ChainDiffVirtualFile(SimpleDiffRequestChain(request), message.filePath)
-                    DiffEditorTabFilesManager.getInstance(context.project).showDiffFile(newDiff, true)
+                val rightDiffContent = if (message.deleted || newFileContent == null) {
+                    EmptyContent()
+                } else {
+                    DiffContentFactory.getInstance().create(newFileContent)
                 }
-            }
 
-            else -> {
-                logger.error { "$FEATURE_NAME: OpenDiff event is received for a conversation that has ${session.sessionState.phase} phase" }
-                messenger.sendError(
-                    tabId = message.tabId,
-                    errMessage = message("amazonqFeatureDev.exception.open_diff_failed"),
-                    retries = 0,
-                    conversationId = session.conversationIdUnsafe
-                )
+                val request = SimpleDiffRequest(message.filePath, leftDiffContent, rightDiffContent, null, null)
+                request.putUserData(DiffUserDataKeys.FORCE_READ_ONLY, true)
+
+                val newDiff = ChainDiffVirtualFile(SimpleDiffRequestChain(request), message.filePath)
+                DiffEditorTabFilesManager.getInstance(context.project).showDiffFile(newDiff, true)
             }
         }
     }
@@ -738,6 +748,7 @@ class DocController(
                 SessionStatePhase.CODEGEN -> {
                     onCodeGeneration(session, message, tabId, mode)
                 }
+
                 else -> null
             }
 

diff --git a/...ware/aws/toolkits/jetbrains/services/amazonqDoc/messages/DocMessagePublisherExtensions.kt b/...ware/aws/toolkits/jetbrains/services/amazonqDoc/messages/DocMessagePublisherExtensions.kt
@@ -7,7 +7,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.ProgressField
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.PromptProgressMessage
-import software.aws.toolkits.jetbrains.services.amazonqDoc.NEW_SESSION_FOLLOWUPS
+import software.aws.toolkits.jetbrains.services.amazonqDoc.ui.NEW_SESSION_FOLLOWUPS
 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

diff --git a/...s-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocSession.kt b/...s-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocSession.kt
@@ -20,7 +20,6 @@ import software.aws.toolkits.jetbrains.common.session.SessionStateConfigData
 import software.aws.toolkits.jetbrains.common.util.AmazonQCodeGenService
 import software.aws.toolkits.jetbrains.common.util.resolveAndCreateOrUpdateFile
 import software.aws.toolkits.jetbrains.common.util.resolveAndDeleteFile
-import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonqDoc.CODE_GENERATION_RETRY_LIMIT
 import software.aws.toolkits.jetbrains.services.amazonqDoc.FEATURE_NAME
@@ -40,7 +39,7 @@ import java.security.MessageDigest
 private val logger = getLogger<AmazonQCodeGenerateClient>()
 
 class DocSession(val tabID: String, val project: Project) {
-    var context: FeatureDevSessionContext
+    var context: DocSessionContext = DocSessionContext(project, MAX_PROJECT_SIZE_BYTES)
     val sessionStartTime = System.currentTimeMillis()
 
     var state: SessionState?
@@ -59,7 +58,6 @@ class DocSession(val tabID: String, val project: Project) {
     var isAuthenticating: Boolean
 
     init {
-        context = FeatureDevSessionContext(project, MAX_PROJECT_SIZE_BYTES)
         proxyClient = AmazonQCodeGenerateClient.getInstance(project)
         amazonQCodeGenService = AmazonQCodeGenService(proxyClient, project)
         state = ConversationNotStartedState("", tabID, token = null)

diff --git a/...nity/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocSessionContext.kt b/...nity/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocSessionContext.kt
@@ -0,0 +1,33 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonqDoc.session
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
+import software.aws.toolkits.jetbrains.services.amazonqDoc.SUPPORTED_DIAGRAM_EXT_SET
+import software.aws.toolkits.jetbrains.services.amazonqDoc.SUPPORTED_DIAGRAM_FILE_NAME_SET
+
+class DocSessionContext(project: Project, maxProjectSizeBytes: Long? = null) : FeatureDevSessionContext(project, maxProjectSizeBytes) {
+
+    /**
+     * Ensure diagram files are not ignored
+     */
+    override fun getAdditionalGitIgnoreBinaryFilesRules(): Set<String> {
+        val ignoreRules = super.getAdditionalGitIgnoreBinaryFilesRules()
+        val diagramExtRulesInGitIgnoreFormatSet = SUPPORTED_DIAGRAM_EXT_SET.map { "*.$it" }.toSet()
+        return ignoreRules - diagramExtRulesInGitIgnoreFormatSet
+    }
+
+    /**
+     * Ensure diagram files are not filtered
+     */
+    override fun isFileExtensionAllowed(file: VirtualFile): Boolean {
+        if (super.isFileExtensionAllowed(file)) {
+            return true
+        }
+
+        return file.extension != null && SUPPORTED_DIAGRAM_FILE_NAME_SET.contains(file.name)
+    }
+}
diff --git a/...brains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/ui/UiContants.kt b/...brains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/ui/UiContants.kt
@@ -0,0 +1,22 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonqDoc.ui
+
+import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUp
+import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUpStatusType
+import software.aws.toolkits.jetbrains.services.amazonqDoc.messages.FollowUpTypes
+import software.aws.toolkits.resources.message
+
+val NEW_SESSION_FOLLOWUPS: List<FollowUp> = listOf(
+    FollowUp(
+        pillText = message("amazonqDoc.prompt.reject.new_task"),
+        type = FollowUpTypes.NEW_TASK,
+        status = FollowUpStatusType.Info
+    ),
+    FollowUp(
+        pillText = message("amazonqDoc.prompt.reject.close_session"),
+        type = FollowUpTypes.CLOSE_SESSION,
+        status = FollowUpStatusType.Info
+    )
+)
diff --git a/...ommunity/src/software/aws/toolkits/jetbrains/services/amazonq/FeatureDevSessionContext.kt b/...ommunity/src/software/aws/toolkits/jetbrains/services/amazonq/FeatureDevSessionContext.kt
@@ -43,7 +43,7 @@ interface RepoSizeError {
 }
 class RepoSizeLimitError(override val message: String) : RuntimeException(), RepoSizeError
 
-class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Long? = null) {
+open class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Long? = null) {
     // TODO: Need to correct this class location in the modules going further to support both amazonq and codescan.
 
     private val additionalGitIgnoreFolderRules = setOf(
@@ -61,7 +61,7 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
         "dist",
     )
 
-    private val additionalGitIgnoreBinaryFilesRules = setOf(
+    private val defaultAdditionalGitIgnoreBinaryFilesRules = setOf(
         "*.zip",
         "*.bin",
         "*.png",
@@ -91,17 +91,17 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
     // 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 var ignorePatternsForBinaryFiles = additionalGitIgnoreBinaryFilesRules
-        .map { convertGitIgnorePatternToRegex(it) }
-        .mapNotNull { pattern ->
-            runCatching { Regex(pattern) }.getOrNull()
-        }
+    private var ignorePatternsForBinaryFiles = buildIgnorePatternsForBinaryFiles()
+
     private val gitIgnoreFile = File(selectedSourceFolder.path, ".gitignore")
 
     init {
         ignorePatternsWithGitIgnore = try {
             buildList {
-                addAll(additionalGitIgnoreFolderRules.map { convertGitIgnorePatternToRegex(it) })
+                addAll(
+                    additionalGitIgnoreFolderRules
+                        .map { convertGitIgnorePatternToRegex(it) }
+                )
                 addAll(parseGitIgnore())
             }.mapNotNull { pattern ->
                 runCatching { Regex(pattern) }.getOrNull()
@@ -111,6 +111,15 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
         }
     }
 
+    private fun buildIgnorePatternsForBinaryFiles(): List<Regex> =
+        getAdditionalGitIgnoreBinaryFilesRules()
+            .map { convertGitIgnorePatternToRegex(it) }
+            .mapNotNull { pattern ->
+                runCatching { Regex(pattern) }.getOrNull()
+            }
+
+    open fun getAdditionalGitIgnoreBinaryFilesRules(): Set<String> = defaultAdditionalGitIgnoreBinaryFilesRules
+
     // This function checks for existence of `devfile.yaml` in customer's repository, currently only `devfile.yaml` is supported for this feature.
     fun checkForDevFile(): Boolean {
         val devFile = File(projectRoot.path, "/devfile.yaml")
@@ -129,7 +138,7 @@ class FeatureDevSessionContext(val project: Project, val maxProjectSizeBytes: Lo
         return ZipCreationResult(zippedProject, checkSum256, zippedProject.length())
     }
 
-    fun isFileExtensionAllowed(file: VirtualFile): Boolean {
+    open fun isFileExtensionAllowed(file: VirtualFile): Boolean {
         // if it is a directory, it is allowed
         if (file.isDirectory) return true
         val extension = file.extension ?: return false