Code transform feature - Human in the loop (#4345)

Author: mk-fan

.changes/next-release/feature-64287d92-9c80-4d31-9e7e-ae1ecdd828ce.json

@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Human in the loop - Adding human intervention to help update dependencies during the Amazon Q Transformation process"
+}

gradle/libs.versions.toml

@@ -22,7 +22,7 @@ mockito = "5.11.0"
 mockitoKotlin = "5.2.1"
 mockk = "1.13.10"
 node-gradle = "7.0.1"
-telemetryGenerator = "1.0.206"
+telemetryGenerator = "1.0.207"
 testLogger = "4.0.0"
 testRetry = "1.5.2"
 # test-only; platform provides slf4j transitively at runtime. <233, 1.7.36; >=233, 2.0.9

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt

@@ -148,6 +148,7 @@ class FeatureDevClient(private val project: Project) {
     suspend fun exportTaskAssistResultArchive(conversationId: String): MutableList<ByteArray> = amazonQStreamingClient.exportResultArchive(
         conversationId,
         ExportIntent.TASK_ASSIST,
+        null,
         { e ->
             LOG.error(e) { "TaskAssist - ExportResultArchive stream exportId=$conversationId exportIntent=${ExportIntent.TASK_ASSIST} Failed: ${e.message} " }
         },

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClientTest.kt

@@ -76,7 +76,7 @@ class AmazonQStreamingClientTest : AmazonQTestBase() {
         val requestCaptor = argumentCaptor<ExportResultArchiveRequest>()
         val handlerCaptor = argumentCaptor<ExportResultArchiveResponseHandler>()
 
-        amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, {}, {})
+        amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, null, {}, {})
         argumentCaptor<ExportResultArchiveRequest, ExportResultArchiveResponseHandler>().apply {
             verify(streamingBearerClient).exportResultArchive(requestCaptor.capture(), handlerCaptor.capture())
         }

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClientTest.kt

@@ -153,11 +153,19 @@ class FeatureDevClientTest : FeatureDevTestBase() {
 
     @Test
     fun `check exportTaskAssistResultArchive`() = runTest {
-        whenever(amazonQStreamingClient.exportResultArchive(any<String>(), any<ExportIntent>(), any(), any())) doReturn exampleExportResultArchiveResponse
+        whenever(
+            amazonQStreamingClient.exportResultArchive(
+                any<String>(),
+                any<ExportIntent>(),
+                eq(null),
+                any(),
+                any()
+            )
+        ) doReturn exampleExportResultArchiveResponse
 
         val actual = featureDevClient.exportTaskAssistResultArchive("1234")
 
-        verify(amazonQStreamingClient).exportResultArchive(eq("1234"), eq(ExportIntent.TASK_ASSIST), any(), any())
+        verify(amazonQStreamingClient).exportResultArchive(eq("1234"), eq(ExportIntent.TASK_ASSIST), eq(null), any(), any())
         verifyNoInteractions(bearerClient)
         verifyNoInteractions(streamingBearerClient)
         verifyNoMoreInteractions(amazonQStreamingClient)

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

@@ -15,25 +15,32 @@ import com.intellij.openapi.vcs.changes.patch.ImportToShelfExecutor
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import software.aws.toolkits.core.utils.error
 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.getCoroutineBgContext
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerArtifact
+import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeTransformHilDownloadArtifact
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.JobId
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.TROUBLESHOOTING_URL_DOWNLOAD_DIFF
+import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getPathToHilArtifactDir
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.openTroubleshootingGuideNotificationAction
 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.CodeTransformApiNames
+import java.io.File
 import java.nio.file.Files
 import java.nio.file.Path
 import java.time.Instant
 import java.util.concurrent.atomic.AtomicBoolean
 
 data class DownloadArtifactResult(val artifact: CodeModernizerArtifact?, val zipPath: String)
+
 class ArtifactHandler(private val project: Project, private val clientAdaptor: GumbyClient) {
     private val telemetry = CodeTransformTelemetryManager.getInstance(project)
     private val downloadedArtifacts = mutableMapOf<JobId, Path>()
@@ -58,6 +65,50 @@ class ArtifactHandler(private val project: Project, private val clientAdaptor: G
         )
     }
 
+    private 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 = try {
+            clientAdaptor.downloadExportResultArchive(jobId, artifactId)
+        } catch (e: Exception) {
+            val errorMessage = "Unexpected error when downloading hil artifact: ${e.localizedMessage}"
+            LOG.error { errorMessage }
+            telemetry.apiError(errorMessage, CodeTransformApiNames.ExportResultArchive, jobId = jobId.id)
+            throw e
+        }
+
+        return try {
+            val tmpPath = tmpDir.toPath()
+            val (downloadZipFilePath, _) = unzipToPath(downloadResultsResponse, tmpPath)
+            LOG.info { "Successfully converted the hil artifact download to a zip at ${downloadZipFilePath.toAbsolutePath()}." }
+            CodeTransformHilDownloadArtifact.create(downloadZipFilePath, getPathToHilArtifactDir(tmpPath))
+        } catch (e: Exception) {
+            // In case if unzip or file operations fail
+            val errorMessage = "Unexpected error when saving downloaded hil artifact: ${e.localizedMessage}"
+            telemetry.error(errorMessage)
+            LOG.error { errorMessage }
+            null
+        }
+    }
+
     suspend fun downloadArtifact(job: JobId): DownloadArtifactResult {
         isCurrentlyDownloading.set(true)
         val downloadStartTime = Instant.now()
@@ -83,16 +134,10 @@ class ArtifactHandler(private val project: Project, private val clientAdaptor: G
 
             // 3. Convert to zip
             LOG.info { "Downloaded the export result archive, about to transform to zip" }
-            val path = Files.createTempFile(null, ".zip")
-            var totalDownloadBytes = 0
-            Files.newOutputStream(path).use {
-                for (bytes in downloadResultsResponse) {
-                    it.write(bytes)
-                    totalDownloadBytes += bytes.size
-                }
-            }
-            LOG.info { "Successfully converted the download to a zip at ${path.toAbsolutePath()}." }
+
+            val (path, totalDownloadBytes) = unzipToPath(downloadResultsResponse)
             val zipPath = path.toAbsolutePath().toString()
+            LOG.info { "Successfully converted the download to a zip at $zipPath." }
 
             // 4. Deserialize zip to CodeModernizerArtifact
             var telemetryErrorMessage: String? = null