suspend fun pollUntilJobCompletion()

in plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt [438:566]
108 lines of code
15 McCabe index (conditional complexity)

    suspend fun pollUntilJobCompletion(
        transformType: CodeTransformType,
        jobId: JobId,
        jobTransitionHandler: (currentStatus: TransformationStatus, migrationPlan: TransformationPlan?) -> Unit,
    ): CodeModernizerJobCompletedResult {
        try {
            state.currentJobId = jobId

            // add delay to avoid the throttling error
            delay(1000)

            var isTransformationPlanEditorOpened = false
            var passedBuild = false
            var passedStart = false

            val result = jobId.pollTransformationStatusAndPlan(
                transformType,
                succeedOn = setOf(
                    TransformationStatus.COMPLETED,
                    TransformationStatus.PAUSED,
                    TransformationStatus.STOPPED,
                    TransformationStatus.PARTIALLY_COMPLETED,
                ),
                failOn = setOf(
                    TransformationStatus.FAILED,
                    TransformationStatus.UNKNOWN_TO_SDK_VERSION,
                ),
                clientAdaptor,
                initialPollingSleepDurationMillis,
                totalPollingSleepDurationMillis,
                isDisposed,
                sessionContext.project,
            ) { old, new, plan ->
                // Always refresh the dev tool tree so status will be up-to-date
                state.currentJobStatus = new
                state.transformationPlan = plan

                if (state.currentJobStatus == TransformationStatus.PAUSED) {
                    val pausedUpdate =
                        state.transformationPlan
                            ?.transformationSteps()
                            ?.flatMap { step -> step.progressUpdates() }
                            ?.filter { update -> update.status() == TransformationProgressUpdateStatus.PAUSED }
                    if (pausedUpdate?.isNotEmpty() == true) {
                        state.currentHilArtifactId = pausedUpdate[0].downloadArtifacts()[0].downloadArtifactId()
                    }
                }

                // Open the transformation plan detail panel once transformation plan is available (no plan for SQL conversions)
                if (transformType != CodeTransformType.SQL_CONVERSION && state.transformationPlan != null && !isTransformationPlanEditorOpened) {
                    tryOpenTransformationPlanEditor()
                    isTransformationPlanEditorOpened = true
                }
                val instant = Instant.now()
                // Set the job start time
                if (state.currentJobCreationTime == Instant.MIN) {
                    state.currentJobCreationTime = instant
                }
                state.updateJobHistory(sessionContext, new, instant)
                setCurrentJobStopTime(new, instant)
                setCurrentJobSummary(new)

                if (!passedStart && new in STATES_AFTER_STARTED) {
                    passedStart = true
                }
                if (!passedBuild && new in STATES_AFTER_INITIAL_BUILD) {
                    passedBuild = true
                }

                jobTransitionHandler(new, plan)
                LOG.info { "Waiting for Modernization Job [$jobId] to complete. State changed for job: $old -> $new" }
            }
            return when {
                result.state == TransformationStatus.STOPPED -> CodeModernizerJobCompletedResult.Stopped

                result.state == TransformationStatus.PAUSED -> CodeModernizerJobCompletedResult.JobPaused(jobId, state.currentHilArtifactId.orEmpty())

                result.state == TransformationStatus.UNKNOWN_TO_SDK_VERSION -> CodeModernizerJobCompletedResult.JobFailed(
                    jobId,
                    message("codemodernizer.notification.warn.unknown_status_response")
                )

                result.state == TransformationStatus.PARTIALLY_COMPLETED -> CodeModernizerJobCompletedResult.JobPartiallySucceeded(jobId)

                result.state == TransformationStatus.FAILED -> {
                    if (!passedStart) {
                        val failureReason = result.jobDetails?.reason() ?: message("codemodernizer.notification.warn.unknown_start_failure")
                        return CodeModernizerJobCompletedResult.JobFailed(jobId, failureReason)
                    } else if (!passedBuild) {
                        // 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 is DownloadArtifactResult.Success) {
                            val failureReason = result.jobDetails?.reason() ?: message("codemodernizer.notification.warn.maven_failed.content")
                            return CodeModernizerJobCompletedResult.JobFailedInitialBuild(jobId, failureReason, true)
                        } else {
                            val failureReason = result.jobDetails?.reason() ?: message("codemodernizer.notification.warn.maven_failed.content")
                            return CodeModernizerJobCompletedResult.JobFailedInitialBuild(jobId, failureReason, false)
                        }
                    } else {
                        val failureReason = result.jobDetails?.reason() ?: message("codemodernizer.notification.warn.unknown_status_response")
                        return CodeModernizerJobCompletedResult.JobFailed(jobId, failureReason)
                    }
                }

                result.state == TransformationStatus.COMPLETED -> {
                    CodeModernizerJobCompletedResult.JobCompletedSuccessfully(jobId)
                }

                // Should not happen
                else -> CodeModernizerJobCompletedResult.JobFailed(jobId, result.jobDetails?.reason().orEmpty())
            }
        } catch (e: Exception) {
            return when (e) {
                is AlreadyDisposedException, is CancellationException -> {
                    LOG.warn { "The session was disposed while polling for job details." }
                    CodeModernizerJobCompletedResult.ManagerDisposed
                }

                else -> {
                    LOG.error(e) { e.message.toString() }
                    CodeModernizerJobCompletedResult.RetryableFailure(
                        jobId,
                        message("codemodernizer.notification.info.modernize_failed.connection_failed", e.message.orEmpty()),
                    )
                }
            }
        }
    }