Q Code Transform: Fix failed uploads and downloads when credentials expire (#4573)

~11% of API failures in May are due to expired credentials being used (InvalidGrantException). The most common places in the code where this happens is when users are about to upload their code or when about to download the code. This is due to at these points customers may have been inactive for a substantial amount of time (>15 min) either due to long running `mvn` builds or due to being afk before downloading results.  

This PR fixes this by forcing a refresh through chat when token is marked as NEED_REFRESH or if we get a SssOidcException. Additionally in the case of downloading an archive we need to resume the job after credential refresh such that users can reattempt the download. Detailed flow for the two cases detailed below:

### Flow when token expired when calling CreateUpload:
1. Chat asks user to reauthenticate
2. Notification asks user to reauthenticate
3. Clicking either causes reauthentication flow to run
4. When user successfully authenticates then the user is requested to start a new flow

### Flow when token expired before calling ExportResultArchive
  1. Chat asks user to reauthenticate
  2. Notification asks user to reauth
  3. Clicking either causes reauthentication flow to run
  4. When user authenticates successfully then the job is resumed
  6. When job fully resumes users are notified to download the results again as normal
     
### Misc changes (bugfixes / supporting changes)

  1. [Supporting change] - Tab opens on resume without user needing to enter /transform
  2. [Supporting change] - Added BearerToken listener to listen for authentication successes so that chat could be informed when user reauthenticates
  3. [Bug fix / Supporting change] - Removed notification “Job completed when you were away” as it was a duplicate and messes up resume flow (as users can click while we are attempting to resume job which causes issues in chat)
  4. [Bug fix / Supporting change] - Removed notification “Job failed when you were away” as it was a duplicate and messes up resume flow (as users can click while we are attempting to resume job which causes issues in chat)
  5. [Bug fix] - Removed notification when job failed during download when more specific error notifications are already shown.

Author: arangatang

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

@@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.codemodernizer
 
 import com.intellij.notification.NotificationAction
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.PersistentStateComponent
@@ -47,6 +48,7 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.model.JobId
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.MAVEN_CONFIGURATION_FILE_NAME
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.MavenCopyCommandsResult
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.MavenDependencyReportCommandsResult
+import software.aws.toolkits.jetbrains.services.codemodernizer.model.UploadFailureReason
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.ValidationResult
 import software.aws.toolkits.jetbrains.services.codemodernizer.panels.managers.CodeModernizerBottomWindowPanelManager
 import software.aws.toolkits.jetbrains.services.codemodernizer.state.CodeModernizerSessionState
@@ -197,6 +199,7 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
      */
     fun addCodeModernizeUI(setSelected: Boolean = false, moduleOrProjectNameForFile: String? = null) = runInEdt {
         val appModernizerBottomWindow = getBottomToolWindow()
+        appModernizerBottomWindow.isAvailable = true
         if (!appModernizerBottomWindow.contentManager.contents.contains(codeModernizerBottomWindowPanelContent)) {
             appModernizerBottomWindow.contentManager.addContent(codeModernizerBottomWindowPanelContent)
         }
@@ -249,13 +252,13 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
         CodeModernizerSessionState.getInstance(project).currentJobStopTime = Instant.MIN
     }
 
-    private fun notifyJobFailure(failureReason: String?) {
+    private fun notifyJobFailure(failureReason: String?, actions: Collection<AnAction> = listOf()) {
         val reason = failureReason ?: message("codemodernizer.notification.info.modernize_failed.unknown_failure_reason") // should not happen
         notifyStickyInfo(
             message("codemodernizer.notification.info.modernize_failed.title"),
             reason,
             project,
-            listOf(displayFeedbackNotificationAction())
+            listOf(displayFeedbackNotificationAction(), *actions.toTypedArray())
         )
     }
 
@@ -301,6 +304,8 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
     }
 
     fun resumeJob(session: CodeModernizerSession, lastJobId: JobId, currentJobResult: TransformationJob) = projectCoroutineScope(project).launch {
+        isJobSuccessfullyResumed.set(true)
+        CodeTransformMessageListener.instance.onTransformResuming()
         if (isModernizationJobActive()) {
             runInEdt { getBottomToolWindow().show() }
             return@launch
@@ -340,6 +345,8 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
         managerState.flags[StateFlags.IS_ONGOING] = false
     }
 
+    fun isJobOngoingInState() = managerState.flags.getOrDefault(StateFlags.IS_ONGOING, false)
+
     fun handleLocalMavenBuildResult(mavenCopyCommandsResult: MavenCopyCommandsResult) {
         codeTransformationSession?.setLastMvnBuildResult(mavenCopyCommandsResult)
         // Send IDE notifications first
@@ -467,6 +474,9 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
         }
     }
 
+    /**
+     * Silently try to resume the job, informs users only when job successfully resumed, suppresses exceptions.
+     */
     fun tryResumeJob(onProjectFirstOpen: Boolean = false) = projectCoroutineScope(project).launch {
         try {
             val notYetResumed = isResumingJob.compareAndSet(false, true)
@@ -492,40 +502,16 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
             val result = session.getJobDetails(lastJobId)
             when (result.status()) {
                 TransformationStatus.COMPLETED -> {
-                    CodeTransformMessageListener.instance.onTransformResuming()
-                    notifyStickyInfo(
-                        message("codemodernizer.manager.job_finished_title"),
-                        message("codemodernizer.manager.job_finished_content"),
-                        project,
-                        listOf(
-                            displayDiffNotificationAction(lastJobId),
-                            displaySummaryNotificationAction(lastJobId),
-                            viewTransformationHubAction()
-                        )
-                    )
                     resumeJob(session, lastJobId, result)
                     setJobNotOngoing()
                 }
 
                 TransformationStatus.PARTIALLY_COMPLETED -> {
-                    CodeTransformMessageListener.instance.onTransformResuming()
-                    notifyStickyInfo(
-                        message("codemodernizer.notification.info.modernize_failed.title"),
-                        message("codemodernizer.manager.job_failed_content", result.reason()),
-                        project,
-                        listOf(
-                            displayDiffNotificationAction(lastJobId),
-                            displaySummaryNotificationAction(lastJobId),
-                            displayFeedbackNotificationAction(),
-                            viewTransformationHubAction()
-                        )
-                    )
                     resumeJob(session, lastJobId, result)
                     setJobNotOngoing()
                 }
 
                 TransformationStatus.UNKNOWN_TO_SDK_VERSION -> {
-                    CodeTransformMessageListener.instance.onTransformResuming()
                     notifyStickyInfo(
                         message("codemodernizer.notification.warn.on_resume.unknown_status_response.title"),
                         message("codemodernizer.notification.warn.on_resume.unknown_status_response.content"),
@@ -534,13 +520,11 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
                 }
 
                 TransformationStatus.STOPPED, TransformationStatus.STOPPING -> {
-                    CodeTransformMessageListener.instance.onTransformResuming()
+                    // If user stopped the last job, there is no need for us to resume the job
                     setJobNotOngoing()
                 }
 
                 else -> {
-                    isJobSuccessfullyResumed.set(true)
-                    CodeTransformMessageListener.instance.onTransformResuming()
                     resumeJob(session, lastJobId, result)
                     notifyStickyInfo(
                         message("codemodernizer.manager.job_ongoing_title"),
@@ -561,10 +545,6 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
         }
     }
 
-    private fun viewTransformationHubAction() = NotificationAction.createSimple(message("codemodernizer.notification.info.modernize_complete.view_summary")) {
-        getBottomToolWindow().show()
-    }
-
     private fun resumeJobNotificationAction(session: CodeModernizerSession, lastJobId: JobId, currentJobResult: TransformationJob) =
         NotificationAction.createSimple(message("codemodernizer.notification.info.modernize_ongoing.view_status")) {
             resumeJob(session, lastJobId, currentJobResult)
@@ -593,9 +573,24 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
                 result.failureReason
             )
 
-            is CodeModernizerJobCompletedResult.ZipUploadFailed -> notifyJobFailure(
-                message("codemodernizer.notification.warn.upload_failed", result.failureReason.toString()),
-            )
+            is CodeModernizerJobCompletedResult.ZipUploadFailed -> {
+                if (result.failureReason is UploadFailureReason.CREDENTIALS_EXPIRED) {
+                    setJobNotOngoing()
+                    CodeTransformMessageListener.instance.onCheckAuth()
+                    notifyJobFailure(
+                        message("codemodernizer.notification.warn.upload_failed_expired_credentials.content"),
+                        listOf(
+                            NotificationAction.createSimpleExpiring(message("codemodernizer.notification.warn.action.reauthenticate")) {
+                                CodeTransformMessageListener.instance.onReauthStarted()
+                            }
+                        )
+                    )
+                } else {
+                    notifyJobFailure(
+                        message("codemodernizer.notification.warn.upload_failed", result.failureReason.toString()),
+                    )
+                }
+            }
 
             is CodeModernizerJobCompletedResult.RetryableFailure -> notifyJobFailure(
                 result.failureReason
@@ -902,4 +897,17 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
 
         return true
     }
+
+    /**
+     * When customer attempts to download an artifact and it fails for some reason (credential expiry etc)
+     * we need to be able to resume the job in order for customers to be able to reattempt the download.
+     * This sets the job as ongoing in the persistent state so that when tryResumeJob is triggered the
+     * IDE attempts to resume the job.
+     */
+    fun handleResumableDownloadArtifactFailure(job: JobId) {
+        // handle the case when user clicks long living notification but has a new job running
+        val session = this.codeTransformationSession ?: return
+        if (session.getActiveJobId() != job) return
+        setJobOngoing(job, session.sessionContext)
+    }
 }

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

@@ -19,6 +19,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Transformation
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationUserActionStatus
 import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
+import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
 import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.exists
 import software.aws.toolkits.core.utils.getLogger
@@ -31,6 +32,7 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModerni
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerSessionContext
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerStartJobResult
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeTransformHilDownloadArtifact
+import software.aws.toolkits.jetbrains.services.codemodernizer.model.DownloadArtifactResult
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.JobId
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.MavenCopyCommandsResult
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.MavenDependencyReportCommandsResult
@@ -42,6 +44,7 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.STATES_AFTE
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.STATES_AFTER_STARTED
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getModuleOrProjectNameForFile
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getPathToHilDependencyReportDir
+import software.aws.toolkits.jetbrains.services.codemodernizer.utils.isValidCodeTransformConnection
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.pollTransformationStatusAndPlan
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.toTransformationLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanSession
@@ -146,6 +149,10 @@ class CodeModernizerSession(
         val payload: File?
 
         // Generate zip file
+        if (!isValidCodeTransformConnection(sessionContext.project)) {
+            // Creating zip can take some time, so quit early
+            return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.CREDENTIALS_EXPIRED)
+        }
         try {
             if (isDisposed.get()) {
                 LOG.warn { "Disposed when about to create zip to upload" }
@@ -181,6 +188,9 @@ class CodeModernizerSession(
         }
 
         // Create upload url and upload zip
+        if (!isValidCodeTransformConnection(sessionContext.project)) {
+            return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.CREDENTIALS_EXPIRED)
+        }
         var uploadId = ""
         try {
             if (shouldStop.get()) {
@@ -218,6 +228,10 @@ class CodeModernizerSession(
                 state.currentJobStatus = TransformationStatus.FAILED
                 return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.OTHER(e.localizedMessage.toString()))
             }
+        } catch (e: SsoOidcException) {
+            state.putJobHistory(sessionContext, TransformationStatus.FAILED)
+            state.currentJobStatus = TransformationStatus.FAILED
+            return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.CREDENTIALS_EXPIRED)
         } catch (e: SdkClientException) {
             // Errors from code whisperer client will always be thrown as SdkClientException
             state.putJobHistory(sessionContext, TransformationStatus.FAILED)
@@ -486,7 +500,7 @@ class CodeModernizerSession(
                         // This is a short term solution to check if build log is available by attempting to download it.
                         // In the long term, we should check if build log is available from transformation metadata.
                         val downloadArtifactResult = artifactHandler.downloadArtifact(jobId, TransformationDownloadArtifactType.LOGS, true)
-                        if (downloadArtifactResult.artifact != null) {
+                        if (downloadArtifactResult is DownloadArtifactResult.Success) {
                             val failureReason = result.jobDetails?.reason() ?: message("codemodernizer.notification.warn.maven_failed.content")
                             return CodeModernizerJobCompletedResult.JobFailedInitialBuild(jobId, failureReason, true)
                         } else {