override suspend fun processButtonClickedMessage()

in plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt [503:1019]


    override suspend fun processButtonClickedMessage(message: IncomingCodeTestMessage.ButtonClicked) {
        val session = codeTestChatHelper.getActiveSession()
        var numberOfLinesGenerated = 0
        var numberOfLinesSelected = 0
        var lineDifference = 0
        var numberOfCharsGenerated = 0
        var numberOfCharsSelected = 0
        var charDifference = 0
        var generatedFileContent = ""
        var selectedFileContent = ""

        when (message.actionID) {
            "utg_view_diff" -> {
                withContext(EDT) {
                    // virtual file only needed for syntax highlighting when viewing diff
                    val tempPath = Files.createTempFile(null, ".${session.testFileName.substringAfterLast('.')}")
                    val virtualFile = tempPath.toFile().toVirtualFile()

                    (DiffManager.getInstance() as DiffManagerEx).showDiffBuiltin(
                        context.project,
                        SimpleDiffRequest(
                            session.testFileName,
                            DiffContentFactory.getInstance().create(
                                getFileContentAtTestFilePath(
                                    session.projectRoot,
                                    session.testFileRelativePathToProjectRoot
                                ),
                                virtualFile
                            ),
                            DiffContentFactory.getInstance().create(
                                session.generatedTestDiffs.values.first(),
                                virtualFile
                            ),
                            "Before",
                            "After"
                        )
                    )
                    Files.deleteIfExists(tempPath)
                    session.openedDiffFile = FileEditorManager.getInstance(context.project).selectedEditor?.file
                    ApplicationManager.getApplication().runReadAction {
                        generatedFileContent = getGeneratedFileContent(session)
                    }

                    selectedFileContent = getFileContentAtTestFilePath(
                        session.projectRoot,
                        session.testFileRelativePathToProjectRoot,
                    )

                    // Line difference calculation: linesOfCodeGenerated = number of lines in generated test file - number of lines in original test file
                    numberOfLinesGenerated = generatedFileContent.lines().size
                    numberOfLinesSelected = selectedFileContent.lines().size
                    lineDifference = numberOfLinesGenerated - numberOfLinesSelected

                    // Character difference calculation: charsOfCodeGenerated = number of characters in generated test file - number of characters in original test file
                    numberOfCharsGenerated = generatedFileContent.length
                    numberOfCharsSelected = selectedFileContent.length
                    charDifference = numberOfCharsGenerated - numberOfCharsSelected

                    session.linesOfCodeGenerated = lineDifference.coerceAtLeast(0)
                    session.charsOfCodeGenerated = charDifference.coerceAtLeast(0)
                    session.latencyOfTestGeneration = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration)
                    UiTelemetry.click(context.project, "unitTestGeneration_viewDiff")

                    val buttonList = mutableListOf<Button>()
                    buttonList.add(
                        Button(
                            "utg_reject",
                            "Reject",
                            keepCardAfterClick = true,
                            position = "outside",
                            status = "error",
                        ),
                    )
                    /*
                    // TODO: for unit test regeneration loop
                    if (session.iteration < 2) {
                        buttonList.add(
                            Button(
                                "utg_regenerate",
                                "Regenerate",
                                keepCardAfterClick = true,
                                position = "outside",
                                status = "info",
                            ),
                        )
                    }
                     */

                    buttonList.add(
                        Button(
                            "utg_accept",
                            "Accept",
                            keepCardAfterClick = true,
                            position = "outside",
                            status = "success",
                        ),
                    )

                    codeTestChatHelper.updateUI(
                        promptInputDisabledState = true,
                        promptInputPlaceholder = message("testgen.placeholder.select_an_option"),
                    )

                    codeTestChatHelper.updateAnswer(
                        CodeTestChatMessageContent(
                            type = ChatMessageType.AnswerPart,
                            buttons = buttonList,
                        ),
                        messageIdOverride = session.viewDiffMessageId
                    )
                }
            }
            "utg_accept" -> {
                // open the file at test path relative to the project root
                val testFileAbsolutePath = Paths.get(session.projectRoot, session.testFileRelativePathToProjectRoot)
                openOrCreateTestFileAndApplyDiff(context.project, testFileAbsolutePath, session.generatedTestDiffs.values.first(), session.openedDiffFile)
                session.codeReferences?.let { references ->
                    LOG.debug { "Accepted unit tests with references: $references" }
                    val manager = CodeWhispererCodeReferenceManager.getInstance(context.project)
                    references.forEach { ref ->
                        var referenceContentSpan: Span? = null
                        ref.recommendationContentSpan()?.let {
                            referenceContentSpan = Span.builder().start(ref.recommendationContentSpan().start())
                                .end(ref.recommendationContentSpan().end()).build()
                        }
                        val reference = Reference.builder().url(
                            ref.url()
                        ).licenseName(ref.licenseName()).repository(ref.repository()).recommendationContentSpan(referenceContentSpan).build()
                        var originalContent: String? = null
                        ref.recommendationContentSpan()?.let {
                            originalContent = session.generatedTestDiffs.values.first().substring(
                                ref.recommendationContentSpan().start(),
                                ref.recommendationContentSpan().end()
                            )
                        }
                        LOG.debug { "Original code content from reference span: $originalContent" }
                        withContext(EDT) {
                            manager.addReferenceLogPanelEntry(reference = reference, null, null, originalContent?.split("\n"))
                            manager.toolWindow?.show()
                        }
                    }
                }
                val testGenerationEventResponse = client.sendTestGenerationEvent(
                    session.testGenerationJob,
                    session.testGenerationJobGroupName,
                    session.programmingLanguage,
                    IdeCategory.JETBRAINS,
                    session.numberOfUnitTestCasesGenerated,
                    session.numberOfUnitTestCasesGenerated,
                    session.linesOfCodeGenerated,
                    session.linesOfCodeGenerated,
                    session.charsOfCodeGenerated,
                    session.charsOfCodeGenerated
                )
                LOG.debug {
                    "Successfully sent test generation telemetry. RequestId: ${
                        testGenerationEventResponse.responseMetadata().requestId()}"
                }

                UiTelemetry.click(context.project, "unitTestGeneration_acceptDiff")

                AmazonqTelemetry.utgGenerateTests(
                    cwsprChatProgrammingLanguage = session.programmingLanguage.languageId,
                    hasUserPromptSupplied = session.hasUserPromptSupplied,
                    isFileInWorkspace = true,
                    isSupportedLanguage = true,
                    credentialStartUrl = getStartUrl(project = context.project),
                    jobGroup = session.testGenerationJobGroupName,
                    jobId = session.testGenerationJob,
                    acceptedCount = session.numberOfUnitTestCasesGenerated?.toLong(),
                    generatedCount = session.numberOfUnitTestCasesGenerated?.toLong(),
                    acceptedLinesCount = session.linesOfCodeGenerated?.toLong(),
                    generatedLinesCount = session.linesOfCodeGenerated?.toLong(),
                    acceptedCharactersCount = session.charsOfCodeGenerated?.toLong(),
                    generatedCharactersCount = session.charsOfCodeGenerated?.toLong(),
                    result = MetricResult.Succeeded,
                    perfClientLatency = session.latencyOfTestGeneration,
                    isCodeBlockSelected = session.isCodeBlockSelected,
                    artifactsUploadDuration = session.artifactUploadDuration,
                    buildPayloadBytes = session.srcPayloadSize,
                    buildZipFileBytes = session.srcZipFileSize,
                    requestId = session.startTestGenerationRequestId,
                    status = Status.ACCEPTED,
                )
                codeTestChatHelper.addAnswer(
                    CodeTestChatMessageContent(
                        message = message("testgen.message.success"),
                        type = ChatMessageType.Answer,
                        canBeVoted = false,
                        buttons = this.showFeedbackButton()
                    )
                )
                codeTestChatHelper.updateUI(
                    promptInputDisabledState = false,
                    promptInputPlaceholder = message("testgen.placeholder.enter_slash_quick_actions"),
                )
                /*
                val taskContext = session.buildAndExecuteTaskContext
                if (session.iteration < 2) {
                    taskContext.buildCommand = getBuildCommand(message.tabId)
                    taskContext.executionCommand = getExecutionCommand(message.tabId)
                    codeTestChatHelper.addAnswer(
                        CodeTestChatMessageContent(
                            message = """
                           Would you like me to help build and execute the test? I'll run following commands

                           ```sh
                           ${taskContext.buildCommand}
                           ${taskContext.executionCommand}
                           ```
                            """.trimIndent(),
                            type = ChatMessageType.Answer,
                            canBeVoted = true,
                            buttons = listOf(
                                Button(
                                    "utg_skip_and_finish",
                                    "Skip and finish",
                                    keepCardAfterClick = true,
                                    position = "outside",
                                    status = "info",
                                ),
                                Button(
                                    "utg_modify_command",
                                    "Modify commands",
                                    keepCardAfterClick = true,
                                    position = "outside",
                                    status = "info",
                                ),
                                Button(
                                    "utg_build_and_execute",
                                    "Build and execute",
                                    keepCardAfterClick = true,
                                    position = "outside",
                                    status = "info",
                                ),
                            )
                        )
                    )
                    codeTestChatHelper.updateUI(
                        promptInputDisabledState = true,
                    )
                } else if (session.iteration < 4) {
                    // Already built and executed once, display # of iterations left message
                    val remainingIterationsCount = UTG_CHAT_MAX_ITERATION - session.iteration
                    val iterationCountString = "$remainingIterationsCount ${if (remainingIterationsCount > 1) "iterations" else "iteration"}"
                    codeTestChatHelper.addAnswer(
                        CodeTestChatMessageContent(
                            message = """
                                    Would you like Amazon Q to build and execute again, and fix errors?

                                    You have $iterationCountString left.

                            """.trimIndent(),
                            type = ChatMessageType.AIPrompt,
                            buttons = listOf(
                                Button(
                                    "utg_skip_and_finish",
                                    "Skip and finish",
                                    keepCardAfterClick = true,
                                    position = "outside",
                                    status = "info",
                                ),
                                Button(
                                    "utg_proceed",
                                    "Proceed",
                                    keepCardAfterClick = true,
                                    position = "outside",
                                    status = "info",
                                ),
                            ),
                        )
                    )
                    codeTestChatHelper.updateUI(
                        promptInputDisabledState = true,
                    )
                } else {
                    // TODO: change this hardcoded string
                    val monthlyLimitString = "25 out of 30"
                    codeTestChatHelper.addAnswer(
                        CodeTestChatMessageContent(
                            message = """
                                 You have gone through all three iterations and this unit test generation workflow is complete. You have $monthlyLimitString Amazon Q Developer Agent invocations left this month.
                            """.trimIndent(),
                            type = ChatMessageType.Answer,
                        )
                    )
                    codeTestChatHelper.updateUI(
                        promptInputPlaceholder = message("testgen.placeholder.newtab")
                    )
                }
                 */
            }
            /*
            //TODO: this is for unit test regeneration build iteration loop
            "utg_regenerate" -> {
                // close the existing open diff in the editor.
                ApplicationManager.getApplication().invokeLater {
                    session.openedDiffFile?.let { FileEditorManager.getInstance(context.project).closeFile(it) }
                }
                codeTestChatHelper.addAnswer(
                    CodeTestChatMessageContent(
                        message = message("testgen.message.regenerate_input"),
                        type = ChatMessageType.Answer,
                        canBeVoted = false
                    )
                )
                val testGenerationEventResponse = client.sendTestGenerationEvent(
                    session.testGenerationJob,
                    session.testGenerationJobGroupName,
                    session.programmingLanguage,
                    session.numberOfUnitTestCasesGenerated,
                    0,
                    session.linesOfCodeGenerated,
                    0,
                    session.charsOfCodeGenerated,
                    0
                )
                LOG.debug {
                    "Successfully sent test generation telemetry. RequestId: ${
                        testGenerationEventResponse.responseMetadata().requestId()}"
                }
                sessionCleanUp(session.tabId)
                codeTestChatHelper.updateUI(
                    promptInputDisabledState = false,
                    promptInputPlaceholder = message("testgen.placeholder.waiting_on_your_inputs"),
                )
            }
             */

            "utg_reject" -> {
                ApplicationManager.getApplication().invokeLater {
                    session.openedDiffFile?.let { FileEditorManager.getInstance(context.project).closeFile(it) }
                }
                codeTestChatHelper.addAnswer(
                    CodeTestChatMessageContent(
                        message = message("testgen.message.success"),
                        type = ChatMessageType.Answer,
                        canBeVoted = false,
                        buttons = this.showFeedbackButton()
                    )
                )
                val testGenerationEventResponse = client.sendTestGenerationEvent(
                    session.testGenerationJob,
                    session.testGenerationJobGroupName,
                    session.programmingLanguage,
                    IdeCategory.JETBRAINS,
                    session.numberOfUnitTestCasesGenerated,
                    0,
                    session.linesOfCodeGenerated,
                    0,
                    session.charsOfCodeGenerated,
                    0
                )
                LOG.debug {
                    "Successfully sent test generation telemetry. RequestId: ${
                        testGenerationEventResponse.responseMetadata().requestId()}"
                }

                UiTelemetry.click(null as Project?, "unitTestGeneration_rejectDiff")
                AmazonqTelemetry.utgGenerateTests(
                    cwsprChatProgrammingLanguage = session.programmingLanguage.languageId,
                    hasUserPromptSupplied = session.hasUserPromptSupplied,
                    isFileInWorkspace = true,
                    isSupportedLanguage = true,
                    credentialStartUrl = getStartUrl(project = context.project),
                    jobGroup = session.testGenerationJobGroupName,
                    jobId = session.testGenerationJob,
                    acceptedCount = 0,
                    generatedCount = session.numberOfUnitTestCasesGenerated?.toLong(),
                    acceptedLinesCount = 0,
                    generatedLinesCount = session.linesOfCodeGenerated?.toLong(),
                    acceptedCharactersCount = 0,
                    generatedCharactersCount = session.charsOfCodeGenerated?.toLong(),
                    result = MetricResult.Succeeded,
                    perfClientLatency = session.latencyOfTestGeneration,
                    isCodeBlockSelected = session.isCodeBlockSelected,
                    artifactsUploadDuration = session.artifactUploadDuration,
                    buildPayloadBytes = session.srcPayloadSize,
                    buildZipFileBytes = session.srcZipFileSize,
                    requestId = session.startTestGenerationRequestId,
                    status = Status.REJECTED,
                )
            }
            "utg_feedback" -> {
                sendFeedback()
                UiTelemetry.click(context.project, "unitTestGeneration_provideFeedback")
            }
            "utg_skip_and_finish" -> {
                codeTestChatHelper.addAnswer(
                    CodeTestChatMessageContent(
                        message = message("testgen.message.success"),
                        type = ChatMessageType.Answer,
                        canBeVoted = false
                    )
                )
                sessionCleanUp(message.tabId)
            }
            "utg_proceed", "utg_build_and_execute" -> {
                // handle both "Proceed" and "Build and execute" button clicks since their actions are similar
                // TODO: show install dependencies card if needed
                session.conversationState = ConversationState.IN_PROGRESS

                // display build in progress card
                val taskContext = session.buildAndExecuteTaskContext

                taskContext.progressStatus = BuildAndExecuteProgressStatus.RUN_BUILD
                val messageId = updateBuildAndExecuteProgressCard(taskContext.progressStatus, null, session.iteration)
                // TODO: build and execute case
                val buildLogsFile = VirtualFileManager.getInstance().findFileByNioPath(
                    withContext(currentCoroutineContext()) {
                        Files.createTempFile(null, null)
                    }
                )
                if (buildLogsFile == null) {
                    // TODO: handle no log file case
                    return
                }
                LOG.debug {
                    "Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
                        "tmpFile for build logs:\n ${buildLogsFile.path}"
                }

                runBuildOrTestCommand(taskContext.buildCommand, buildLogsFile, context.project, isBuildCommand = true, taskContext)
                while (taskContext.buildExitCode < 0) {
                    // wait until build command finished
                    delay(1000)
                }

                // TODO: only go to future iterations when buildExitCode or testExitCode > 0, right now iterate regardless
                if (taskContext.buildExitCode > 0) {
                    // TODO: handle build failure case
                    // ...
//                    return
                }
                taskContext.progressStatus = BuildAndExecuteProgressStatus.RUN_EXECUTION_TESTS
                updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)

                val testLogsFile = VirtualFileManager.getInstance().findFileByNioPath(
                    withContext(currentCoroutineContext()) {
                        Files.createTempFile(null, null)
                    }
                )
                if (testLogsFile == null) {
                    // TODO: handle no log file case
                    return
                }
                LOG.debug {
                    "Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
                        "tmpFile for test logs:\n ${buildLogsFile.path}"
                }
                delay(1000)
                runBuildOrTestCommand(taskContext.executionCommand, testLogsFile, context.project, isBuildCommand = false, taskContext)
                while (taskContext.testExitCode < 0) {
                    // wait until test command finished
                    delay(1000)
                }

                if (taskContext.testExitCode == 0) {
                    taskContext.progressStatus = BuildAndExecuteProgressStatus.TESTS_EXECUTED
                    updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
                    codeTestChatHelper.addAnswer(
                        CodeTestChatMessageContent(
                            message = message("testgen.message.success"),
                            type = ChatMessageType.Answer,
                            canBeVoted = false
                        )
                    )
                    sessionCleanUp(message.tabId)
                    return
                }

                // has test failure, we will zip the latest project and invoke backend again
                taskContext.progressStatus = BuildAndExecuteProgressStatus.FIXING_TEST_CASES
                val buildAndExecuteMessageId = updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)

                val previousUTGIterationContext = PreviousUTGIterationContext(
                    buildLogFile = buildLogsFile,
                    testLogFile = testLogsFile,
                    selectedFile = session.selectedFile,
                    buildAndExecuteMessageId = buildAndExecuteMessageId
                )

                val job = CodeWhispererUTGChatManager.getInstance(context.project).generateTests("", codeTestChatHelper, previousUTGIterationContext, null)
                job?.join()

                taskContext.progressStatus = BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS
                // session.iteration already updated in generateTests
                updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration - 1)
            }
            "utg_modify_command" -> {
                // TODO allow user input to modify the command
                codeTestChatHelper.addAnswer(
                    CodeTestChatMessageContent(
                        message = """
                            Sure. Let me know which command you'd like to modify or you could also provide all command lines you'd like me to run.
                            
                        """.trimIndent(),
                        type = ChatMessageType.Answer,
                        canBeVoted = false
                    )
                )
                session.conversationState = ConversationState.WAITING_FOR_BUILD_COMMAND_INPUT
            }
            "utg_install_and_continue" -> {
                // TODO: install dependencies and build
            }
            "stop_test_generation" -> {
                UiTelemetry.click(null as Project?, "unitTestGeneration_cancelTestGenerationProgress")
                session.isGeneratingTests = false
                sessionCleanUp(message.tabId)
                return
            }
            else -> {
                // Handle other cases or do nothing
            }
        }
    }