Merge branch 'feature/q-lsp-chat' into stopResp

Author: leigaol

plugins/amazonq/codetransform/jetbrains-community/resources/META-INF/codetransform-ext-java.xml

@@ -18,7 +18,7 @@
         <projectService serviceImplementation="software.aws.toolkits.jetbrains.services.codemodernizer.panels.managers.CodeModernizerBottomWindowPanelManager"/>
         <toolWindow id="aws.codewhisperer.codetransform" anchor="bottom" doNotActivateOnStart="true" canCloseContents="false"
                     factoryClass="software.aws.toolkits.jetbrains.services.codemodernizer.toolwindow.CodeModernizerBottomToolWindowFactory"
-                    icon="AllIcons.Actions.Preview"/>
+                    icon="AllIcons.Actions.Properties"/>
         <fileEditorProvider implementation="software.aws.toolkits.jetbrains.services.codemodernizer.plan.CodeModernizerPlanEditorProvider"/>
         <fileEditorProvider implementation="software.aws.toolkits.jetbrains.services.codemodernizer.summary.CodeModernizerSummaryEditorProvider"/>
     </extensions>

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ArtifactHandler.kt

@@ -24,7 +24,6 @@ import software.aws.toolkits.core.utils.exists
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.jetbrains.core.coroutines.EDT
-import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.NoTokenInitializedException
 import software.aws.toolkits.jetbrains.services.amazonq.CODE_TRANSFORM_TROUBLESHOOT_DOC_DOWNLOAD_ERROR_OVERVIEW
@@ -50,13 +49,13 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.state.CodeModerni
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getPathToHilArtifactDir
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.isValidCodeTransformConnection
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.openTroubleshootingGuideNotificationAction
+import software.aws.toolkits.jetbrains.services.codemodernizer.utils.zipToPath
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.jetbrains.utils.notifyStickyInfo
 import software.aws.toolkits.jetbrains.utils.notifyStickyWarn
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodeTransformArtifactType
 import java.io.File
-import java.nio.file.Files
 import java.nio.file.Path
 import java.time.Instant
 import java.util.concurrent.atomic.AtomicBoolean
@@ -112,32 +111,12 @@ class ArtifactHandler(
         }
     }
 
-    suspend fun unzipToPath(byteArrayList: List<ByteArray>, outputDirPath: Path? = null): Pair<Path, Int> {
-        val zipFilePath = withContext(getCoroutineBgContext()) {
-            if (outputDirPath == null) {
-                Files.createTempFile(null, ".zip")
-            } else {
-                Files.createTempFile(outputDirPath, null, ".zip")
-            }
-        }
-        var totalDownloadBytes = 0
-        withContext(getCoroutineBgContext()) {
-            Files.newOutputStream(zipFilePath).use {
-                for (bytes in byteArrayList) {
-                    it.write(bytes)
-                    totalDownloadBytes += bytes.size
-                }
-            }
-        }
-        return zipFilePath to totalDownloadBytes
-    }
-
     suspend fun downloadHilArtifact(jobId: JobId, artifactId: String, tmpDir: File): CodeTransformHilDownloadArtifact? {
         val downloadResultsResponse = clientAdaptor.downloadExportResultArchive(jobId, artifactId)
 
         return try {
             val tmpPath = tmpDir.toPath()
-            val (downloadZipFilePath, _) = unzipToPath(downloadResultsResponse, tmpPath)
+            val (downloadZipFilePath, _) = zipToPath(downloadResultsResponse, tmpPath)
             LOG.info { "Successfully converted the hil artifact download to a zip at ${downloadZipFilePath.toAbsolutePath()}." }
             CodeTransformHilDownloadArtifact.create(downloadZipFilePath, getPathToHilArtifactDir(tmpPath))
         } catch (e: Exception) {
@@ -211,7 +190,7 @@ class ArtifactHandler(
             val totalDownloadBytes: Int
             val zipPath: String
             try {
-                val result = unzipToPath(downloadResultsResponse)
+                val result = zipToPath(downloadResultsResponse)
                 path = result.first
                 totalDownloadBytes = result.second
                 zipPath = path.toAbsolutePath().toString()

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerManager.kt

@@ -909,7 +909,7 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
                 // Add delay between upload complete and trying to resume
                 delay(500)
 
-                codeTransformationSession?.resumeTransformFromHil()
+                codeTransformationSession?.resumeTransformation()
             } else {
                 throw CodeModernizerException("Cannot create dependency zip for HIL")
             }

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt

@@ -20,6 +20,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Transformation
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationProgressUpdateStatus
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationUserActionStatus
+import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
 import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
 import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
 import software.aws.toolkits.core.utils.Waiters.waitUntil
@@ -174,6 +175,7 @@ class CodeModernizerSession(
             }
             // for language upgrades, copyResult should always be Successful here, failure cases already handled
             val result = sessionContext.createZipWithModuleFiles(copyResult)
+            sessionContext.originalUploadZipPath = result.payload.toPath()
 
             if (result is ZipCreationResult.Missing1P) {
                 telemetryErrorMessage = "Missing 1p dependencies"
@@ -283,9 +285,7 @@ class CodeModernizerSession(
             return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.OTHER(e.localizedMessage))
         } finally {
             telemetry.uploadProject(payloadSize, startTime, true, telemetryErrorMessage)
-            if (payload != null) {
-                deleteUploadArtifact(payload)
-            }
+            // do not delete upload ZIP; re-used for client-side build
         }
 
         // Send upload completion message to chat (only if successful)
@@ -309,12 +309,6 @@ class CodeModernizerSession(
         }
     }
 
-    internal fun deleteUploadArtifact(payload: File) {
-        if (!payload.delete()) {
-            LOG.warn { "Unable to delete upload artifact." }
-        }
-    }
-
     private fun startJob(uploadId: String): StartTransformationResponse {
         val sourceLanguage = sessionContext.sourceJavaVersion.name.toTransformationLanguage()
         val targetLanguage = sessionContext.targetJavaVersion.name.toTransformationLanguage()
@@ -341,9 +335,10 @@ class CodeModernizerSession(
      */
     fun resumeJob(startTime: Instant, jobId: JobId) = state.putJobHistory(sessionContext, TransformationStatus.STARTED, jobId.id, startTime)
 
-    fun resumeTransformFromHil() {
+    fun resumeTransformation() {
         val clientAdaptor = GumbyClient.getInstance(sessionContext.project)
         clientAdaptor.resumeCodeTransformation(state.currentJobId as JobId, TransformationUserActionStatus.COMPLETED)
+        getLogger<CodeModernizerManager>().info { "Successfully resumed transformation with status of COMPLETED" }
     }
 
     fun rejectHilAndContinue(): ResumeTransformationResponse {
@@ -390,7 +385,7 @@ class CodeModernizerSession(
     /**
      * Adapted from [CodeWhispererCodeScanSession]
      */
-    suspend fun uploadPayload(payload: File): String {
+    suspend fun uploadPayload(payload: File, uploadContext: UploadContext? = null): String {
         val sha256checksum: String = Base64.getEncoder().encodeToString(
             withContext(getCoroutineBgContext()) {
                 DigestUtils.sha256(FileInputStream(payload))
@@ -400,7 +395,7 @@ class CodeModernizerSession(
             throw AlreadyDisposedException("Disposed when about to create upload URL")
         }
         val clientAdaptor = GumbyClient.getInstance(sessionContext.project)
-        val createUploadUrlResponse = clientAdaptor.createGumbyUploadUrl(sha256checksum)
+        val createUploadUrlResponse = clientAdaptor.createGumbyUploadUrl(sha256checksum, uploadContext)
 
         LOG.info {
             "Uploading project artifact at ${payload.path} with checksum $sha256checksum using uploadId: ${
@@ -428,9 +423,9 @@ class CodeModernizerSession(
                 createUploadUrlResponse.kmsKeyArn().orEmpty(),
             ) { shouldStop.get() }
         }
-        LOG.info { "Upload to S3 succeeded" }
+        LOG.info { "Upload of ${payload.path} to S3 succeeded with upload context of $uploadContext" }
         if (!shouldStop.get()) {
-            LOG.info { "Uploaded artifact. Latency: ${calculateTotalLatency(uploadStartTime, Instant.now())}ms" }
+            LOG.info { "Uploaded artifact. Latency: ${calculateTotalLatency(uploadStartTime, Instant.now())} ms" }
         }
         return createUploadUrlResponse.uploadId()
     }

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeTransformChatApp.kt

@@ -41,9 +41,11 @@ private enum class CodeTransformMessageTypes(val type: String) {
     ChatPrompt("chat-prompt"), // for getting the transformation objective
     CodeTransformStart("codetransform-start"),
     CodeTransformSelectSQLMetadata("codetransform-select-sql-metadata"),
+    CodeTransformConfirmCustomDependencyVersions("codetransform-input-confirm-custom-dependency-versions"),
     CodeTransformSelectSQLModuleSchema("codetransform-select-sql-module-schema"),
     CodeTransformStop("codetransform-stop"),
     CodeTransformCancel("codetransform-cancel"),
+    CodeTransformContinue("codetransform-continue"),
     CodeTransformConfirmSkipTests("codetransform-confirm-skip-tests"),
     CodeTransformConfirmOneOrMultipleDiffs("codetransform-confirm-one-or-multiple-diffs"),
     CodeTransformNew("codetransform-new"),
@@ -74,9 +76,12 @@ class CodeTransformChatApp : AmazonQApp {
             CodeTransformMessageTypes.Transform.type to IncomingCodeTransformMessage.Transform::class,
             CodeTransformMessageTypes.CodeTransformStart.type to IncomingCodeTransformMessage.CodeTransformStart::class,
             CodeTransformMessageTypes.CodeTransformSelectSQLMetadata.type to IncomingCodeTransformMessage.CodeTransformSelectSQLMetadata::class,
+            CodeTransformMessageTypes.CodeTransformConfirmCustomDependencyVersions.type to
+                IncomingCodeTransformMessage.CodeTransformConfirmCustomDependencyVersions::class,
             CodeTransformMessageTypes.CodeTransformSelectSQLModuleSchema.type to IncomingCodeTransformMessage.CodeTransformSelectSQLModuleSchema::class,
             CodeTransformMessageTypes.CodeTransformStop.type to IncomingCodeTransformMessage.CodeTransformStop::class,
             CodeTransformMessageTypes.CodeTransformCancel.type to IncomingCodeTransformMessage.CodeTransformCancel::class,
+            CodeTransformMessageTypes.CodeTransformContinue.type to IncomingCodeTransformMessage.CodeTransformContinue::class,
             CodeTransformMessageTypes.ChatPrompt.type to IncomingCodeTransformMessage.ChatPrompt::class,
             CodeTransformMessageTypes.CodeTransformConfirmSkipTests.type to IncomingCodeTransformMessage.CodeTransformConfirmSkipTests::class,
             CodeTransformMessageTypes.CodeTransformConfirmOneOrMultipleDiffs.type to IncomingCodeTransformMessage.CodeTransformConfirmOneOrMultipleDiffs::class,
@@ -182,12 +187,14 @@ class CodeTransformChatApp : AmazonQApp {
             is IncomingCodeTransformMessage.CodeTransformSelectSQLMetadata -> inboundAppMessagesHandler.processCodeTransformSelectSQLMetadataAction(message)
             is IncomingCodeTransformMessage.CodeTransformSelectSQLModuleSchema ->
                 inboundAppMessagesHandler.processCodeTransformSelectSQLModuleSchemaAction(message)
-
+            is IncomingCodeTransformMessage.CodeTransformConfirmCustomDependencyVersions ->
+                inboundAppMessagesHandler.processCodeTransformCustomDependencyVersions(message)
             is IncomingCodeTransformMessage.CodeTransformCancel -> inboundAppMessagesHandler.processCodeTransformCancelAction(message)
             is IncomingCodeTransformMessage.CodeTransformStop -> inboundAppMessagesHandler.processCodeTransformStopAction(message.tabId)
             is IncomingCodeTransformMessage.ChatPrompt -> inboundAppMessagesHandler.processChatPromptMessage(message)
             is IncomingCodeTransformMessage.CodeTransformConfirmSkipTests -> inboundAppMessagesHandler.processCodeTransformConfirmSkipTests(message)
             is IncomingCodeTransformMessage.CodeTransformConfirmOneOrMultipleDiffs -> inboundAppMessagesHandler.processCodeTransformOneOrMultipleDiffs(message)
+            is IncomingCodeTransformMessage.CodeTransformContinue -> inboundAppMessagesHandler.processCodeTransformContinueAction(message)
             is IncomingCodeTransformMessage.CodeTransformNew -> inboundAppMessagesHandler.processCodeTransformNewAction(message)
             is IncomingCodeTransformMessage.CodeTransformOpenTransformHub -> inboundAppMessagesHandler.processCodeTransformOpenTransformHub(message)
             is IncomingCodeTransformMessage.CodeTransformOpenMvnBuild -> inboundAppMessagesHandler.processCodeTransformOpenMvnBuild(message)

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/InboundAppMessagesHandler.kt

@@ -11,12 +11,16 @@ interface InboundAppMessagesHandler {
 
     suspend fun processCodeTransformCancelAction(message: IncomingCodeTransformMessage.CodeTransformCancel)
 
+    suspend fun processCodeTransformContinueAction(message: IncomingCodeTransformMessage.CodeTransformContinue)
+
     suspend fun processCodeTransformStartAction(message: IncomingCodeTransformMessage.CodeTransformStart)
 
     suspend fun processCodeTransformSelectSQLMetadataAction(message: IncomingCodeTransformMessage.CodeTransformSelectSQLMetadata)
 
     suspend fun processCodeTransformSelectSQLModuleSchemaAction(message: IncomingCodeTransformMessage.CodeTransformSelectSQLModuleSchema)
 
+    suspend fun processCodeTransformCustomDependencyVersions(message: IncomingCodeTransformMessage.CodeTransformConfirmCustomDependencyVersions)
+
     suspend fun processCodeTransformStopAction(tabId: String)
 
     suspend fun processChatPromptMessage(message: IncomingCodeTransformMessage.ChatPrompt)

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt

@@ -61,12 +61,13 @@ class GumbyClient(private val project: Project) {
     private val amazonQStreamingClient
         get() = AmazonQStreamingClient.getInstance(project)
 
-    fun createGumbyUploadUrl(sha256Checksum: String): CreateUploadUrlResponse {
+    fun createGumbyUploadUrl(sha256Checksum: String, context: UploadContext? = null): CreateUploadUrlResponse {
         val request = CreateUploadUrlRequest.builder()
             .contentChecksumType(ContentChecksumType.SHA_256)
             .contentChecksum(sha256Checksum)
             .uploadIntent(UploadIntent.TRANSFORMATION)
             .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn)
+            .uploadContext(context)
             .build()
         return callApi({ bearerClient().createUploadUrl(request) }, apiName = "CreateUploadUrl")
     }
@@ -164,27 +165,27 @@ class GumbyClient(private val project: Project) {
 
     suspend fun downloadExportResultArchive(
         jobId: JobId,
-        hilDownloadArtifactId: String? = null,
+        downloadArtifactId: String? = null,
         downloadArtifactType: TransformationDownloadArtifactType? = TransformationDownloadArtifactType.CLIENT_INSTRUCTIONS,
     ): MutableList<ByteArray> = amazonQStreamingClient.exportResultArchive(
         jobId.id,
         ExportIntent.TRANSFORMATION,
-        if (hilDownloadArtifactId == null) {
+        if (downloadArtifactId == null) {
             null
         } else {
             ExportContext
                 .builder()
                 .transformationExportContext(
                     TransformationExportContext
                         .builder()
-                        .downloadArtifactId(hilDownloadArtifactId)
+                        .downloadArtifactId(downloadArtifactId)
                         .downloadArtifactType(downloadArtifactType.toString())
                         .build()
                 )
                 .build()
         },
         { e ->
-            LOG.error(e) { "ExportResultArchive failed: ${e.message}" }
+            LOG.error(e) { "ExportResultArchive on job ${jobId.id} and artifact $downloadArtifactId failed: ${e.message}" }
         },
         { startTime ->
             LOG.info { "ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" }