From 6fb2dda3f528d863fc2eb5409e0e7865fc93a32d Mon Sep 17 00:00:00 2001 From: Laxman Reddy <141967714+laileni-aws@users.noreply.github.com> Date: Thu, 23 Oct 2025 14:16:32 -0700 Subject: [PATCH 1/4] refactor(amazonq): removing /agents code (#6078) * Refactor: removing /test code * Refactor: removing /doc code --- .../clients/AmazonQCodeGenerateClient.kt | 189 --- .../session/ConversationNotStartedState.kt | 21 - .../jetbrains/common/session/SessionState.kt | 17 - .../common/session/SessionStateTypes.kt | 24 - .../common/util/AmazonQCodeGenService.kt | 248 --- .../amazonqCodeTest/CodeTestChatApp.kt | 106 -- .../amazonqCodeTest/CodeTestChatAppFactory.kt | 11 - .../amazonqCodeTest/CodeTestChatItems.kt | 38 - .../amazonqCodeTest/CodeTestConstants.kt | 19 - .../CodeWhispererCodeTestSession.kt | 146 -- .../CodeWhispererUTGChatManager.kt | 589 ------- .../InboundAppMessagesHandler.kt | 30 - .../controller/CodeTestChatController.kt | 1445 ----------------- .../controller/CodeTestChatHelper.kt | 182 --- .../messages/CodeTestMessage.kt | 257 --- .../model/BuildAndExecuteStatusIcon.kt | 39 - .../model/PreviousUTGIterationContext.kt | 13 - .../amazonqCodeTest/model/ShortAnswer.kt | 41 - .../amazonqCodeTest/session/Session.kt | 72 - .../storage/ChatSessionStorage.kt | 24 - .../amazonqCodeTest/utils/UTGChatUtil.kt | 200 --- .../jetbrains/services/amazonqDoc/DocApp.kt | 108 -- .../services/amazonqDoc/DocAppFactory.kt | 11 - .../services/amazonqDoc/DocChatItems.kt | 45 - .../services/amazonqDoc/DocConstants.kt | 41 - .../services/amazonqDoc/DocExceptions.kt | 95 -- .../amazonqDoc/InboundAppMessagesHandler.kt | 18 - .../amazonqDoc/controller/DocController.kt | 1090 ------------- .../controller/DocControllerExtensions.kt | 156 -- .../controller/DocGenerationTask.kt | 99 -- .../amazonqDoc/messages/DocMessage.kt | 296 ---- .../messages/DocMessagePublisherExtensions.kt | 260 --- .../amazonqDoc/session/DocGenerationState.kt | 292 ---- .../services/amazonqDoc/session/DocSession.kt | 341 ---- .../amazonqDoc/session/DocSessionContext.kt | 9 - .../session/PrepareDocGenerationState.kt | 82 - .../amazonqDoc/session/SessionStateTypes.kt | 27 - .../amazonqDoc/storage/ChatSessionStorage.kt | 33 - .../services/amazonqDoc/ui/UiContants.kt | 22 - .../amazonqDoc/util/DocControllerUtil.kt | 43 - 40 files changed, 6779 deletions(-) delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/ConversationNotStartedState.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionState.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionStateTypes.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/util/AmazonQCodeGenService.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatAppFactory.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatItems.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestConstants.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/InboundAppMessagesHandler.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatHelper.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/messages/CodeTestMessage.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/model/BuildAndExecuteStatusIcon.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/model/PreviousUTGIterationContext.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/model/ShortAnswer.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/session/Session.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/storage/ChatSessionStorage.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/utils/UTGChatUtil.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocApp.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocAppFactory.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocChatItems.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocConstants.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocExceptions.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/InboundAppMessagesHandler.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocController.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocControllerExtensions.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/controller/DocGenerationTask.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/messages/DocMessage.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/messages/DocMessagePublisherExtensions.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocGenerationState.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocSession.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/DocSessionContext.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/PrepareDocGenerationState.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/session/SessionStateTypes.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/storage/ChatSessionStorage.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/ui/UiContants.kt delete mode 100644 plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/util/DocControllerUtil.kt diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt deleted file mode 100644 index 3957d9fdf46..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.common.clients - -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.project.Project -import com.intellij.openapi.util.SystemInfo -import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient -import software.amazon.awssdk.services.codewhispererruntime.model.ArtifactType -import software.amazon.awssdk.services.codewhispererruntime.model.ContentChecksumType -import software.amazon.awssdk.services.codewhispererruntime.model.CreateTaskAssistConversationRequest -import software.amazon.awssdk.services.codewhispererruntime.model.CreateTaskAssistConversationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse -import software.amazon.awssdk.services.codewhispererruntime.model.Dimension -import software.amazon.awssdk.services.codewhispererruntime.model.DocV2AcceptanceEvent -import software.amazon.awssdk.services.codewhispererruntime.model.DocV2GenerationEvent -import software.amazon.awssdk.services.codewhispererruntime.model.GetTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.IdeCategory -import software.amazon.awssdk.services.codewhispererruntime.model.OperatingSystem -import software.amazon.awssdk.services.codewhispererruntime.model.OptOutPreference -import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse -import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.TaskAssistPlanningUploadContext -import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext -import software.amazon.awssdk.services.codewhispererruntime.model.UploadIntent -import software.amazon.awssdk.services.codewhispererruntime.model.UserContext -import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent -import software.aws.toolkits.core.utils.error -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.core.utils.info -import software.aws.toolkits.jetbrains.common.session.Intent -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection -import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager -import software.aws.toolkits.jetbrains.services.amazonqDoc.FEATURE_EVALUATION_PRODUCT_NAME -import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency -import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata -import software.aws.toolkits.jetbrains.settings.AwsSettings -import java.time.Instant -import software.amazon.awssdk.services.codewhispererruntime.model.ChatTriggerType as SyncChatTriggerType - -@Service(Service.Level.PROJECT) -class AmazonQCodeGenerateClient(private val project: Project) { - private fun getTelemetryOptOutPreference() = - if (AwsSettings.getInstance().isTelemetryEnabled) { - OptOutPreference.OPTIN - } else { - OptOutPreference.OPTOUT - } - - private val docUserContext = ClientMetadata.getDefault().let { - val osForFeatureDev: OperatingSystem = - when { - SystemInfo.isWindows -> OperatingSystem.WINDOWS - SystemInfo.isMac -> OperatingSystem.MAC - // For now, categorize everything else as "Linux" (Linux/FreeBSD/Solaris/etc.) - else -> OperatingSystem.LINUX - } - - UserContext.builder() - .ideCategory(IdeCategory.JETBRAINS) - .operatingSystem(osForFeatureDev) - .product(FEATURE_EVALUATION_PRODUCT_NAME) - .clientId(it.clientId) - .ideVersion(it.awsVersion) - .build() - } - - fun connection() = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - ?: error("Attempted to use connection while one does not exist") - - fun bearerClient() = QRegionProfileManager.getInstance().getQClient(project) - - private val amazonQStreamingClient - get() = AmazonQStreamingClient.getInstance(project) - - fun sendDocTelemetryEvent( - generationEvent: DocV2GenerationEvent? = null, - acceptanceEvent: DocV2AcceptanceEvent? = null, - ): SendTelemetryEventResponse = - bearerClient().sendTelemetryEvent { requestBuilder -> - requestBuilder.telemetryEvent { telemetryEventBuilder -> - generationEvent?.let { telemetryEventBuilder.docV2GenerationEvent(it) } - acceptanceEvent?.let { telemetryEventBuilder.docV2AcceptanceEvent(it) } - } - requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(docUserContext) - requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun sendDocMetricData(operationName: String, result: String): SendTelemetryEventResponse = - bearerClient().sendTelemetryEvent { requestBuilder -> - requestBuilder.telemetryEvent { telemetryEventBuilder -> - telemetryEventBuilder.metricData { - it - .metricName("Operation") - .metricValue(1.0) - .timestamp(Instant.now()) - .product("DocGeneration") - .dimensions( - listOf( - Dimension.builder() - .name("operationName") - .value(operationName) - .build(), - Dimension.builder() - .name("result") - .value(result) - .build() - ) - ) - } - } - requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(docUserContext) - } - - fun createTaskAssistConversation(): CreateTaskAssistConversationResponse = bearerClient().createTaskAssistConversation( - CreateTaskAssistConversationRequest.builder() - .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - .build() - ) - - fun createTaskAssistUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long): CreateUploadUrlResponse = - bearerClient().createUploadUrl { - it.contentChecksumType(ContentChecksumType.SHA_256) - .contentChecksum(contentChecksumSha256) - .contentLength(contentLength) - .artifactType(ArtifactType.SOURCE_CODE) - .uploadIntent(UploadIntent.TASK_ASSIST_PLANNING) - .uploadContext( - UploadContext.builder() - .taskAssistPlanningUploadContext( - TaskAssistPlanningUploadContext.builder() - .conversationId(conversationId) - .build() - ) - .build() - ) - .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, userMessage: String, intent: Intent): StartTaskAssistCodeGenerationResponse = - bearerClient() - .startTaskAssistCodeGeneration { request -> - request - .conversationState { - it - .conversationId(conversationId) - .chatTriggerType(SyncChatTriggerType.MANUAL) - .currentMessage { cm -> cm.userInputMessage { um -> um.content(userMessage) } } - } - .workspaceState { - it - .programmingLanguage { pl -> pl.languageName("javascript") } - .uploadId(uploadId) - } - .intent(intent.name) - .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun getTaskAssistCodeGeneration(conversationId: String, codeGenerationId: String): GetTaskAssistCodeGenerationResponse = bearerClient() - .getTaskAssistCodeGeneration { - it - .conversationId(conversationId) - .codeGenerationId(codeGenerationId) - } - - suspend fun exportTaskAssistResultArchive(conversationId: String): MutableList = amazonQStreamingClient.exportResultArchive( - conversationId, - ExportIntent.TASK_ASSIST, - null, - { e -> - LOG.error(e) { "TaskAssist - ExportResultArchive stream exportId=$conversationId exportIntent=${ExportIntent.TASK_ASSIST} Failed: ${e.message} " } - }, - { startTime -> - LOG.info { "TaskAssist - ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" } - } - ) - - companion object { - private val LOG = getLogger() - - fun getInstance(project: Project) = project.service() - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/ConversationNotStartedState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/ConversationNotStartedState.kt deleted file mode 100644 index ad300449dfb..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/ConversationNotStartedState.kt +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.common.session - -import software.aws.toolkits.jetbrains.services.amazonqDoc.session.SessionStateInteraction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStateAction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource - -class ConversationNotStartedState( - override var approach: String, - override val tabID: String, - override var token: CancellationTokenSource?, -) : SessionState { - override val phase = SessionStatePhase.INIT - - override suspend fun interact(action: SessionStateAction): SessionStateInteraction { - error("Illegal transition between states, restart the conversation") - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionState.kt deleted file mode 100644 index 15850963820..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionState.kt +++ /dev/null @@ -1,17 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.common.session - -import software.aws.toolkits.jetbrains.services.amazonqDoc.session.SessionStateInteraction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStateAction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource - -interface SessionState { - val tabID: String - val phase: SessionStatePhase? - var token: CancellationTokenSource? - var approach: String - suspend fun interact(action: SessionStateAction): SessionStateInteraction -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionStateTypes.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionStateTypes.kt deleted file mode 100644 index 3a75a029a77..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/session/SessionStateTypes.kt +++ /dev/null @@ -1,24 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.common.session - -import software.aws.toolkits.jetbrains.common.util.AmazonQCodeGenService -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext - -open class SessionStateConfig( - open val conversationId: String, - open val repoContext: FeatureDevSessionContext, - open val amazonQCodeGenService: AmazonQCodeGenService, -) - -data class SessionStateConfigData( - override val conversationId: String, - override val repoContext: FeatureDevSessionContext, - override val amazonQCodeGenService: AmazonQCodeGenService, -) : SessionStateConfig(conversationId, repoContext, amazonQCodeGenService) - -enum class Intent { - DEV, - DOC, -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/util/AmazonQCodeGenService.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/util/AmazonQCodeGenService.kt deleted file mode 100644 index 299e9c5ba57..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/util/AmazonQCodeGenService.kt +++ /dev/null @@ -1,248 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.common.util - -import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper -import com.fasterxml.jackson.module.kotlin.readValue -import com.intellij.openapi.project.Project -import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse -import software.amazon.awssdk.services.codewhispererruntime.model.GetTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.ServiceQuotaExceededException -import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException -import software.amazon.awssdk.services.codewhispererruntime.model.ValidationException -import software.amazon.awssdk.services.codewhispererstreaming.model.CodeWhispererStreamingException -import software.aws.toolkits.core.utils.debug -import software.aws.toolkits.core.utils.error -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.common.clients.AmazonQCodeGenerateClient -import software.aws.toolkits.jetbrains.common.session.Intent -import software.aws.toolkits.jetbrains.services.amazonqDoc.session.DocGenerationStreamResult -import software.aws.toolkits.jetbrains.services.amazonqDoc.session.ExportDocTaskAssistResultArchiveStreamResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ApiException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeIterationLimitException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ContentLengthException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ExportParseException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevOperation -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MonthlyConversationLimitError -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ServiceException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ZipFileCorruptedException -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl -import software.aws.toolkits.telemetry.AmazonqTelemetry -import software.aws.toolkits.telemetry.MetricResult -import software.aws.toolkits.jetbrains.services.amazonqDoc.ContentLengthException as DocContentLengthException - -class AmazonQCodeGenService(val proxyClient: AmazonQCodeGenerateClient, val project: Project) { - fun createConversation(): String { - val startTime = System.currentTimeMillis() - var failureReason: String? = null - var failureReasonDesc: String? = null - var result: MetricResult = MetricResult.Succeeded - var conversationId: String? = null - try { - logger.debug { "Executing createTaskAssistConversation" } - val taskAssistConversationResult = proxyClient.createTaskAssistConversation() - conversationId = taskAssistConversationResult.conversationId() - logger.debug { - "$FEATURE_NAME: Created conversation: {conversationId: $conversationId, requestId: ${ - taskAssistConversationResult.responseMetadata().requestId() - }" - } - - return conversationId - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: Failed to start conversation: ${e.message}" } - result = MetricResult.Failed - failureReason = e.javaClass.simpleName - if (e is FeatureDevException) { - failureReason = e.reason() - failureReasonDesc = e.reasonDesc() - } - var errMssg = e.message - if (e is CodeWhispererRuntimeException) { - errMssg = e.awsErrorDetails().errorMessage() - logger.warn(e) { "Start conversation failed for request: ${e.requestId()}" } - - if (e is ServiceQuotaExceededException) { - throw MonthlyConversationLimitError(errMssg, operation = FeatureDevOperation.CreateConversation.toString(), desc = null, cause = e.cause) - } - throw ApiException.of(e.statusCode(), errMssg, operation = FeatureDevOperation.CreateConversation.toString(), desc = null, e.cause) - } - throw ServiceException( - errMssg ?: "CreateTaskAssistConversation failed", - operation = FeatureDevOperation.CreateConversation.toString(), - desc = null, - e.cause - ) - } finally { - AmazonqTelemetry.startConversationInvoke( - amazonqConversationId = conversationId, - result = result, - reason = failureReason, - reasonDesc = failureReasonDesc, - duration = (System.currentTimeMillis() - startTime).toDouble(), - credentialStartUrl = getStartUrl(project = this.project), - ) - } - } - - fun createUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long, uploadId: String, featureName: String? = null): - CreateUploadUrlResponse { - try { - logger.debug { "Executing createUploadUrl with conversationId $conversationId" } - val uploadUrlResponse = proxyClient.createTaskAssistUploadUrl( - conversationId, - contentChecksumSha256, - contentLength, - ) - logger.debug { - "$FEATURE_NAME: Created upload url: {uploadId: $uploadId, requestId: ${uploadUrlResponse.responseMetadata().requestId()}}" - } - return uploadUrlResponse - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: Failed to generate presigned url: ${e.message}" } - - var errMssg = e.message - if (e is CodeWhispererRuntimeException) { - errMssg = e.awsErrorDetails().errorMessage() - logger.warn(e) { "Create UploadUrl failed for request: ${e.requestId()}" } - - if (e is ValidationException && e.message?.contains("Invalid contentLength") == true) { - if (featureName?.equals("docGeneration") == true) { - throw DocContentLengthException(operation = FeatureDevOperation.CreateUploadUrl.toString(), desc = null, cause = e.cause) - } - throw ContentLengthException(operation = FeatureDevOperation.CreateUploadUrl.toString(), desc = null, cause = e.cause) - } - - throw ApiException.of(e.statusCode(), errMssg, operation = FeatureDevOperation.CreateUploadUrl.toString(), desc = null, e.cause) - } - throw ServiceException(errMssg ?: "CreateUploadUrl failed", operation = FeatureDevOperation.CreateUploadUrl.toString(), desc = null, e.cause) - } - } - - open fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, message: String, intent: Intent): - StartTaskAssistCodeGenerationResponse { - try { - logger.debug { "Executing startTaskAssistCodeGeneration with conversationId: $conversationId , uploadId: $uploadId" } - val startCodeGenerationResponse = proxyClient.startTaskAssistCodeGeneration( - conversationId, - uploadId, - message, - intent - ) - - logger.debug { "$FEATURE_NAME: Started code generation with requestId: ${startCodeGenerationResponse.responseMetadata().requestId()}" } - return startCodeGenerationResponse - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: Failed to execute startTaskAssistCodeGeneration ${e.message}" } - - var errMssg = e.message - if (e is CodeWhispererRuntimeException) { - errMssg = e.awsErrorDetails().errorMessage() - logger.warn(e) { "StartTaskAssistCodeGeneration failed for request: ${e.requestId()}" } - - // API Front-end will throw Throttling if conversation limit is reached. API Front-end monitors StartCodeGeneration for throttling - if (e is software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException && - e.message?.contains("StartTaskAssistCodeGeneration reached for this month.") == true - ) { - throw MonthlyConversationLimitError(errMssg, operation = FeatureDevOperation.StartTaskAssistCodeGeneration.toString(), desc = null, e.cause) - } - - if (e is ServiceQuotaExceededException || ( - e is ThrottlingException && e.message?.contains( - "limit for number of iterations on a code generation" - ) == true - ) - ) { - throw CodeIterationLimitException(operation = FeatureDevOperation.StartTaskAssistCodeGeneration.toString(), desc = null, e.cause) - } else if (e is ValidationException && e.message?.contains("repo size is exceeding the limits") == true) { - if (intent == Intent.DOC) { - throw DocContentLengthException(operation = FeatureDevOperation.CreateUploadUrl.toString(), desc = null, cause = e.cause) - } - throw ContentLengthException(operation = FeatureDevOperation.StartTaskAssistCodeGeneration.toString(), desc = null, cause = e.cause) - } else if (e is ValidationException && e.message?.contains("zipped file is corrupted") == true) { - throw ZipFileCorruptedException(operation = FeatureDevOperation.StartTaskAssistCodeGeneration.toString(), desc = null, e.cause) - } - throw ApiException.of(e.statusCode(), errMssg, operation = FeatureDevOperation.StartTaskAssistCodeGeneration.toString(), desc = null, e.cause) - } - throw ServiceException( - errMssg ?: "StartTaskAssistCodeGeneration failed", - operation = FeatureDevOperation.StartTaskAssistCodeGeneration.toString(), - desc = null, - e.cause - ) - } - } - - fun getTaskAssistCodeGeneration(conversationId: String, codeGenerationId: String): GetTaskAssistCodeGenerationResponse { - try { - logger.debug { "Executing GetTaskAssistCodeGeneration with conversationId: $conversationId , codeGenerationId: $codeGenerationId" } - val getCodeGenerationResponse = proxyClient.getTaskAssistCodeGeneration(conversationId, codeGenerationId) - - logger.debug { - "$FEATURE_NAME: Received code generation status $getCodeGenerationResponse with requestId ${ - getCodeGenerationResponse.responseMetadata() - .requestId() - }" - } - return getCodeGenerationResponse - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: Failed to execute GetTaskAssistCodeGeneration ${e.message}" } - - var errMssg = e.message - if (e is CodeWhispererRuntimeException) { - errMssg = e.awsErrorDetails().errorMessage() - logger.warn(e) { "GetTaskAssistCodeGeneration failed for request: ${e.requestId()}" } - throw ApiException.of(e.statusCode(), errMssg, operation = FeatureDevOperation.GetTaskAssistCodeGeneration.toString(), desc = null, e.cause) - } - throw ServiceException( - errMssg ?: "GetTaskAssistCodeGeneration failed", - operation = FeatureDevOperation.GetTaskAssistCodeGeneration.toString(), - desc = null, - e.cause - ) - } - } - - suspend fun exportTaskAssistArchiveResult(conversationId: String): DocGenerationStreamResult { - val exportResponse: MutableList - try { - exportResponse = proxyClient.exportTaskAssistResultArchive(conversationId) - logger.debug { "$FEATURE_NAME: Received export task assist result archive response" } - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: Failed to export archive result: ${e.message}" } - - var errMssg = e.message - if (e is CodeWhispererStreamingException) { - errMssg = e.awsErrorDetails().errorMessage() - logger.warn(e) { "ExportTaskAssistArchiveResult failed for request: ${e.requestId()}" } - } - throw ServiceException( - errMssg ?: "ExportTaskAssistArchive failed", - operation = FeatureDevOperation.ExportTaskAssistArchiveResult.toString(), - desc = null, - e.cause - ) - } - - val parsedResult: ExportDocTaskAssistResultArchiveStreamResult - try { - val result = exportResponse.reduce { acc, next -> acc + next } // To map the result it is needed to combine the full byte array - parsedResult = jacksonObjectMapper().readValue(result) - } catch (e: Exception) { - logger.error(e) { "Failed to parse downloaded code results" } - throw ExportParseException(operation = FeatureDevOperation.ExportTaskAssistArchiveResult.toString(), desc = null, e.cause) - } - - return parsedResult.codeGenerationResult - } - - companion object { - private val logger = getLogger() - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt deleted file mode 100644 index 7972e45fb9c..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt +++ /dev/null @@ -1,106 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqCodeTest - -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.project.Project -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.launch -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener -import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp -import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext -import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener -import software.aws.toolkits.jetbrains.services.amazonqCodeScan.auth.isCodeScanAvailable -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.auth.isCodeTestAvailable -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.CodeTestChatController -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.AuthenticationUpdateMessage -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.IncomingCodeTestMessage -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.storage.ChatSessionStorage -import software.aws.toolkits.jetbrains.services.amazonqDoc.auth.isDocAvailable -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.auth.isFeatureDevAvailable -import software.aws.toolkits.jetbrains.services.codemodernizer.utils.isCodeTransformAvailable - -class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp { - - override val tabTypes = listOf("codetest") - - override fun init(context: AmazonQAppInitContext) { - val chatSessionStorage = ChatSessionStorage() - val inboundAppMessagesHandler = - CodeTestChatController(context, chatSessionStorage, cs = scope) - - context.messageTypeRegistry.register( - "clear" to IncomingCodeTestMessage.ClearChat::class, - "help" to IncomingCodeTestMessage.Help::class, - "chat-prompt" to IncomingCodeTestMessage.ChatPrompt::class, - "new-tab-was-created" to IncomingCodeTestMessage.NewTabCreated::class, - "tab-was-removed" to IncomingCodeTestMessage.TabRemoved::class, - "start-test-gen" to IncomingCodeTestMessage.StartTestGen::class, - "response-body-link-click" to IncomingCodeTestMessage.ClickedLink::class, - "button-click" to IncomingCodeTestMessage.ButtonClicked::class, - "chat-item-voted" to IncomingCodeTestMessage.ChatItemVoted::class, - "chat-item-feedback" to IncomingCodeTestMessage.ChatItemFeedback::class, - "button-click" to IncomingCodeTestMessage.ButtonClicked::class, - "auth-follow-up-was-clicked" to IncomingCodeTestMessage.AuthFollowUpWasClicked::class - ) - - scope.launch { - context.messagesFromUiToApp.flow.collect { message -> - // Launch a new coroutine to handle each message - scope.launch { handleMessage(message, inboundAppMessagesHandler) } - } - } - - ApplicationManager.getApplication().messageBus.connect(this).subscribe( - ToolkitConnectionManagerListener.TOPIC, - object : ToolkitConnectionManagerListener { - override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - scope.launch { - context.messagesFromAppToUi.publish( - AuthenticationUpdateMessage( - featureDevEnabled = isFeatureDevAvailable(context.project), - codeTransformEnabled = isCodeTransformAvailable(context.project), - codeScanEnabled = isCodeScanAvailable(context.project), - codeTestEnabled = isCodeTestAvailable(context.project), - docEnabled = isDocAvailable(context.project), - authenticatingTabIDs = chatSessionStorage.getAuthenticatingSessions().map { it.tabId } - ) - ) - } - } - } - ) - - context.project.messageBus.connect(this).subscribe( - QRegionProfileSelectedListener.TOPIC, - object : QRegionProfileSelectedListener { - override fun onProfileSelected(project: Project, profile: QRegionProfile?) { - chatSessionStorage.deleteAllSessions() - } - } - ) - } - - private suspend fun handleMessage(message: AmazonQMessage, inboundAppMessagesHandler: InboundAppMessagesHandler) { - when (message) { - is IncomingCodeTestMessage.ClearChat -> inboundAppMessagesHandler.processClearQuickAction(message) - is IncomingCodeTestMessage.Help -> inboundAppMessagesHandler.processHelpQuickAction(message) - is IncomingCodeTestMessage.ChatPrompt -> inboundAppMessagesHandler.processPromptChatMessage(message) - is IncomingCodeTestMessage.NewTabCreated -> inboundAppMessagesHandler.processNewTabCreatedMessage(message) - is IncomingCodeTestMessage.TabRemoved -> inboundAppMessagesHandler.processTabRemovedMessage(message) - is IncomingCodeTestMessage.StartTestGen -> inboundAppMessagesHandler.processStartTestGen(message) - is IncomingCodeTestMessage.ClickedLink -> inboundAppMessagesHandler.processLinkClick(message) - is IncomingCodeTestMessage.ButtonClicked -> inboundAppMessagesHandler.processButtonClickedMessage(message) - is IncomingCodeTestMessage.ChatItemVoted -> inboundAppMessagesHandler.processChatItemVoted(message) - is IncomingCodeTestMessage.ChatItemFeedback -> inboundAppMessagesHandler.processChatItemFeedBack(message) - is IncomingCodeTestMessage.AuthFollowUpWasClicked -> inboundAppMessagesHandler.processAuthFollowUpClick(message) - } - } - - override fun dispose() { - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatAppFactory.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatAppFactory.kt deleted file mode 100644 index 46ccd461d9a..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatAppFactory.kt +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqCodeTest -import com.intellij.openapi.project.Project -import kotlinx.coroutines.CoroutineScope -import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppFactory - -class CodeTestChatAppFactory(private val cs: CoroutineScope) : AmazonQAppFactory { - override fun createApp(project: Project) = CodeTestChatApp(cs) -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatItems.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatItems.kt deleted file mode 100644 index 4ed35e35065..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatItems.kt +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqCodeTest - -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.Button -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestButtonId -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.ProgressField -import software.aws.toolkits.resources.message - -val cancellingProgressField = ProgressField( - status = "warning", - text = message("general.canceling"), - value = -1, - actions = emptyList() -) - -// TODO: Need to change the string after the F2F -val testGenCompletedField = ProgressField( - status = "success", - text = message("general.success"), - value = 100, - actions = emptyList() -) - -val cancelTestGenButton = Button( - id = CodeTestButtonId.StopTestGeneration.id, - text = message("general.cancel"), - icon = "cancel" -) - -fun testGenProgressField(value: Int) = ProgressField( - status = "default", - text = message("testgen.progressbar.generate_unit_tests"), - value = value, - valueText = "$value%", - actions = listOf(cancelTestGenButton) -) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestConstants.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestConstants.kt deleted file mode 100644 index e5e25ed7cb0..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestConstants.kt +++ /dev/null @@ -1,19 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqCodeTest - -const val FEATURE_NAME = "Amazon Q Unit Test Generation" - -enum class ConversationState { - IDLE, - WAITING_FOR_BUILD_COMMAND_INPUT, - WAITING_FOR_REGENERATE_INPUT, - IN_PROGRESS, -} - -fun generateSummaryMessage(fileName: String): String = """ - Sure. This may take a few minutes. I'll share updates here as I work on this. - **Generating unit tests for the following methods in $fileName**: - -""".trimIndent() diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt deleted file mode 100644 index f13d7fb7f43..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt +++ /dev/null @@ -1,146 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqCodeTest -import com.intellij.openapi.project.Project -import kotlinx.coroutines.ensureActive -import kotlinx.coroutines.time.withTimeout -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse -import software.aws.toolkits.core.utils.debug -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.CodeTestChatHelper -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessageContent -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PreviousUTGIterationContext -import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext -import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig -import software.aws.toolkits.jetbrains.services.codewhisperer.model.CreateUploadUrlServiceInvocationContext -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_BYTES_IN_KB -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_MILLIS_IN_SECOND -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererZipUploadManager -import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType -import software.aws.toolkits.jetbrains.utils.assertIsNonDispatchThread -import java.nio.file.Path -import java.time.Duration -import java.time.Instant -import java.util.UUID -import kotlin.coroutines.coroutineContext - -// TODO: Refactor with CodeWhispererCodeScanSession code since both are about zip CreateUploadUrl logic -class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) { - private fun now() = Instant.now().toEpochMilli() - - /** - * Run UTG sessions are follow steps: - * 1. Zipping project - * 2. Creating Upload url & Upload to S3 bucket - */ - suspend fun run(codeTestChatHelper: CodeTestChatHelper, previousIterationContext: PreviousUTGIterationContext?): CodeTestResponseContext { - try { - assertIsNonDispatchThread() - coroutineContext.ensureActive() - - val path = sessionContext.sessionConfig.getRelativePath() - ?: throw RuntimeException("Can not determine current file path for adding unit tests") - - // Add card answer to show UTG in progress - val testSummaryMessageId = - if (previousIterationContext == null) { - codeTestChatHelper.addAnswer( - CodeTestChatMessageContent( - message = generateSummaryMessage(path.fileName.toString()), - type = ChatMessageType.AnswerStream - ) - ).also { - // For streaming effect - codeTestChatHelper.updateAnswer( - CodeTestChatMessageContent(type = ChatMessageType.AnswerPart) - ) - codeTestChatHelper.updateUI( - loadingChat = true, - promptInputDisabledState = true - ) - if (it == null) { - throw RuntimeException("Can not add test summary card") - } - } - } else { - // non-first iteration doesn't have a test summary card - null - } - - val (payloadContext, sourceZip) = withTimeout(Duration.ofSeconds(sessionContext.sessionConfig.createPayloadTimeoutInSeconds())) { - sessionContext.sessionConfig.createPayload() - } - - LOG.debug { - "Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" + - "Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" + - "Total number of lines included: ${payloadContext.totalLines} \n" + - "Total number of files included in payload: ${payloadContext.totalFiles} \n" + - "Total time taken for creating payload: ${payloadContext.totalTimeInMilliseconds * 1.0 / TOTAL_MILLIS_IN_SECOND} seconds\n" + - "Payload context language: ${payloadContext.language}" + - "Payload exceeded the limit: ${payloadContext.payloadLimitCrossed}" - } - - // 2 & 3. CreateUploadURL and upload the context. - val artifactsUploadStartTime = now() - val taskName = UUID.randomUUID().toString() - val sourceZipUploadResponse = - CodeWhispererZipUploadManager.getInstance(sessionContext.project).createUploadUrlAndUpload( - sourceZip, - "SourceCode", - CodeWhispererConstants.UploadTaskType.UTG, - taskName, - CodeWhispererConstants.FeatureName.TEST_GENERATION - ) - - sourceZipUploadResponse.uploadId() - - LOG.debug { - "Successfully uploaded source zip to s3: " + - "Upload id: ${sourceZipUploadResponse.uploadId()} " + - "Request id: ${sourceZipUploadResponse.responseMetadata().requestId()}" - } - val artifactsUploadDuration = now() - artifactsUploadStartTime - - val codeTestResponseContext = CodeTestResponseContext( - payloadContext, - CreateUploadUrlServiceInvocationContext(artifactsUploadDuration = artifactsUploadDuration), - path, - sourceZipUploadResponse, - testSummaryMessageId - ) - - return codeTestResponseContext - } catch (e: Exception) { - LOG.debug(e) { "Error while creating zip and uploading to S3" } - throw e - } - } - - companion object { - private val LOG = getLogger() - } -} - -sealed class CodeTestResponse { - abstract val responseContext: CodeTestResponseContext - data class Success(val message: String, override val responseContext: CodeTestResponseContext) : CodeTestResponse() - - data class Error(val errorMessage: String, override val responseContext: CodeTestResponseContext) : CodeTestResponse() -} - -data class CodeTestSessionContext( - val project: Project, - val sessionConfig: CodeTestSessionConfig, -) - -data class CodeTestResponseContext( - val payloadContext: PayloadContext, - val serviceInvocationContext: CreateUploadUrlServiceInvocationContext, - val currentFileRelativePath: Path, - val createUploadUrlResponse: CreateUploadUrlResponse, - val testSummaryMessageId: String?, - val reason: String? = null, -) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt deleted file mode 100644 index aecede14c15..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt +++ /dev/null @@ -1,589 +0,0 @@ -// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqCodeTest - -import com.fasterxml.jackson.core.JsonParseException -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.fileEditor.FileDocumentManager -import com.intellij.openapi.fileEditor.FileEditorManager -import com.intellij.openapi.fileEditor.impl.FileDocumentManagerImpl -import com.intellij.openapi.project.Project -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Job -import kotlinx.coroutines.delay -import kotlinx.coroutines.launch -import software.amazon.awssdk.core.exception.SdkServiceException -import software.amazon.awssdk.services.codewhispererruntime.model.GetTestGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.PackageInfo -import software.amazon.awssdk.services.codewhispererruntime.model.Range -import software.amazon.awssdk.services.codewhispererruntime.model.StartTestGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode -import software.amazon.awssdk.services.codewhispererruntime.model.TargetFileInfo -import software.amazon.awssdk.services.codewhispererruntime.model.TestGenerationJobStatus -import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext -import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent -import software.aws.toolkits.core.utils.Waiters.waitUntil -import software.aws.toolkits.core.utils.debug -import software.aws.toolkits.core.utils.error -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.core.utils.info -import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser -import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.CodeTestChatHelper -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.Button -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessageContent -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PreviousUTGIterationContext -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.BuildAndExecuteProgressStatus -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session -import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles -import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency -import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException -import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.fileTooLarge -import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig -import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.testGenStoppedError -import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth -import software.aws.toolkits.jetbrains.services.codewhisperer.util.getTelemetryErrorMessage -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl -import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType -import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference -import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings -import software.aws.toolkits.jetbrains.utils.isQConnected -import software.aws.toolkits.resources.message -import software.aws.toolkits.telemetry.AmazonqTelemetry -import software.aws.toolkits.telemetry.MetricResult -import software.aws.toolkits.telemetry.Status -import java.io.ByteArrayInputStream -import java.io.ByteArrayOutputStream -import java.io.File -import java.io.IOException -import java.nio.file.Paths -import java.time.Duration -import java.time.Instant -import java.util.concurrent.atomic.AtomicBoolean -import java.util.zip.ZipInputStream - -@Service -class CodeWhispererUTGChatManager(val project: Project, private val cs: CoroutineScope) { - // TODO: consider combining this with session.isGeneratingTests - private val isUTGInProgress = AtomicBoolean(false) - private val generatedTestDiffs = mutableMapOf() - - private fun throwIfCancelled(session: Session) { - if (!session.isGeneratingTests) { - testGenStoppedError() - } - } - - private suspend fun launchTestGenFlow( - prompt: String, - codeTestChatHelper: CodeTestChatHelper, - previousIterationContext: PreviousUTGIterationContext?, - selectionRange: Range?, - ) { - // 1st API call: Zip project and call CreateUploadUrl - val session = codeTestChatHelper.getActiveSession() - session.isGeneratingTests = true - session.iteration++ - - // Set progress bar to "Generating unit tests..." - codeTestChatHelper.updateUI( - promptInputDisabledState = true, - promptInputProgress = testGenProgressField(0), - ) - - val codeTestResponseContext = createUploadUrl(codeTestChatHelper, previousIterationContext) - session.srcPayloadSize = codeTestResponseContext.payloadContext.srcPayloadSize - session.srcZipFileSize = codeTestResponseContext.payloadContext.srcZipFileSize - session.artifactUploadDuration = codeTestResponseContext.serviceInvocationContext.artifactsUploadDuration - val path = codeTestResponseContext.currentFileRelativePath - - if (codeTestResponseContext.payloadContext.payloadLimitCrossed == true) { - fileTooLarge() - } - - val createUploadUrlResponse = codeTestResponseContext.createUploadUrlResponse ?: return - throwIfCancelled(session) - - LOG.debug { - "Q TestGen StartTestGenerationRequest: TabId=${codeTestChatHelper.getActiveCodeTestTabId()}, " + - "uploadId=${createUploadUrlResponse.uploadId()}, relativeTargetPath=${codeTestResponseContext.currentFileRelativePath}, " + - "selectionRange=$selectionRange" - } - - // 2nd API call: StartTestGeneration - val startTestGenerationResponse = try { - var response: StartTestGenerationResponse? = null - waitUntil( - succeedOn = { response?.sdkHttpResponse()?.statusCode() == 200 }, - maxDuration = Duration.ofSeconds(1), // 1-second timeout - ) { - response = startTestGeneration( - uploadId = createUploadUrlResponse.uploadId(), - targetCode = listOf( - TargetCode.builder() - .relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString()) - .targetLineRangeList(selectionRange?.let { listOf(it) }.orEmpty()) - .build() - ), - userInput = prompt - ) - delay(200) - response?.testGenerationJob() != null - } - response ?: throw RuntimeException("Failed to start test generation") - } catch (e: Exception) { - LOG.error(e) { "Unexpected error while creating test generation job" } - if (e is SdkServiceException) { - session.startTestGenerationRequestId = e.requestId() - } - throw CodeTestException( - "CreateTestJobError: ${getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)}", - "CreateTestJobError", - message("testgen.error.generic_technical_error_message") - ) - } - - val job = startTestGenerationResponse.testGenerationJob() - session.startTestGenerationRequestId = startTestGenerationResponse.responseMetadata().requestId() - session.testGenerationJobGroupName = job.testGenerationJobGroupName() - session.testGenerationJob = job.testGenerationJobId() - throwIfCancelled(session) - - // 3rd API call: Polling mechanism on test job status - var finished = false - var testGenerationResponse: GetTestGenerationResponse? = null - var packageInfoList = emptyList() - var packageInfo: PackageInfo? = null - var targetFileInfo: TargetFileInfo? = null - - while (!finished) { - throwIfCancelled(session) - testGenerationResponse = getTestGenerationStatus(job.testGenerationJobId(), job.testGenerationJobGroupName()) - val status = testGenerationResponse.testGenerationJob().status() - packageInfoList = testGenerationResponse.testGenerationJob().packageInfoList() - packageInfo = packageInfoList.firstOrNull() - targetFileInfo = packageInfo?.targetFileInfoList()?.firstOrNull() - val filePlan = targetFileInfo?.filePlan() - val cleanedfilePlan = filePlan - ?.replace(Regex("^```\\s*"), "") // Remove leading triple backticks - ?.replace(Regex("\\s*```$"), "") // Remove trailing triple backticks - ?.trim() - - when (status) { - TestGenerationJobStatus.COMPLETED -> { - LOG.debug { "Test generation completed, package info: $packageInfoList" } - finished = true - if (packageInfo != null && targetFileInfo != null) { - session.packageInfoList = packageInfoList - session.testFileName = targetFileInfo.testFilePath()?.let { File(it).name }.orEmpty() - session.numberOfUnitTestCasesGenerated = targetFileInfo.numberOfTestMethods() ?: 0 - session.testFileRelativePathToProjectRoot = getTestFilePathRelativeToRoot(targetFileInfo) - - if (previousIterationContext == null) { - codeTestChatHelper.updateAnswer( - CodeTestChatMessageContent( - message = generateSummaryMessage(path.fileName.toString()) + (cleanedfilePlan.orEmpty()), - type = ChatMessageType.Answer, - ), - messageIdOverride = codeTestResponseContext.testSummaryMessageId - ) - } - } - } - TestGenerationJobStatus.FAILED -> { - LOG.debug { "Test generation failed, package info: $packageInfoList" } - if (testGenerationResponse.testGenerationJob().jobStatusReason() == null) { - throw CodeTestException( - "TestGenFailedError: " + message("testgen.message.failed"), - "TestGenFailedError", - message("testgen.error.generic_technical_error_message") - ) - } else { - throw CodeTestException( - "TestGenFailedError: ${testGenerationResponse.testGenerationJob().jobStatusReason()}", - "TestGenFailedError", - testGenerationResponse.testGenerationJob().jobStatusReason() - ) - } - } - else -> { - LOG.debug { "Test generation in progress, progress rate: ${testGenerationResponse.testGenerationJob().progressRate()}" } - - if (previousIterationContext == null) { - codeTestChatHelper.updateAnswer( - CodeTestChatMessageContent( - message = generateSummaryMessage(path.fileName.toString()) + (cleanedfilePlan.orEmpty()), - type = ChatMessageType.Answer, - ), - messageIdOverride = codeTestResponseContext.testSummaryMessageId - ) - } - codeTestChatHelper.updateUI( - promptInputDisabledState = true, - promptInputProgress = testGenProgressField(testGenerationResponse.testGenerationJob().progressRate() ?: 0), - ) - } - } - // polling every 2 seconds to reduce # of API calls - delay(2000) - } - - throwIfCancelled(session) - - // 4th API call: Step 4: ExportResultsArchive - val byteArray = AmazonQStreamingClient.getInstance(project).exportResultArchive( - createUploadUrlResponse.uploadId(), - ExportIntent.UNIT_TESTS, - ExportContext.fromUnitTestGenerationExportContext { - it.testGenerationJobId(job.testGenerationJobId()) - it.testGenerationJobGroupName(job.testGenerationJobGroupName()) - }, - { e -> - LOG.error(e) { "ExportResultArchive failed: ${e.message}" } - throw CodeTestException( - "ExportResultsArchiveError: ${e.message}", - "ExportResultsArchiveError", - message("testgen.error.generic_technical_error_message") - ) - }, - { startTime -> - LOG.info { "ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" } - } - ) - val result = byteArray.reduce { acc, next -> acc + next } // To map the result it is needed to combine the full byte array - storeGeneratedTestDiffs(result, session) - if (!session.isGeneratingTests) { - // TODO: Modify text according to FnF - codeTestChatHelper.addAnswer( - CodeTestChatMessageContent( - message = message("testgen.error.generic_technical_error_message"), - type = ChatMessageType.Answer, - canBeVoted = true - ) - ) - return - } - - targetFileInfo = session.packageInfoList.firstOrNull()?.targetFileInfoList()?.firstOrNull() - val codeReference = targetFileInfo?.codeReferences()?.map { ref -> - CodeReference( - licenseName = ref.licenseName(), - url = ref.url(), - information = "${ref.licenseName()} - ${ref.repository()}" - ) - } - targetFileInfo?.codeReferences()?.let { session.codeReferences = it } - val isReferenceAllowed = CodeWhispererSettings.getInstance().isIncludeCodeWithReference() - if (!isReferenceAllowed && codeReference?.isNotEmpty() == true) { - codeTestChatHelper.addAnswer( - CodeTestChatMessageContent( - message = """ - Your settings do not allow code generation with references. - """.trimIndent(), - type = ChatMessageType.Answer, - ) - ) - } else { - if (previousIterationContext == null) { - // show another card as the answer, remove this once we remove 3 backticks from backend - val jobSummary = testGenerationResponse?.testGenerationJob()?.jobSummary()?.trim().orEmpty() - - val cleanedPlanSummary = jobSummary - .replace(Regex("(?m)^\\s*```(?:\\w+)?\\s*"), "") // Removes leading backticks with possible language specifier - .replace(Regex("(?m)\\s*```\\s*$"), "") // Removes trailing backticks - .trim() - - val summary = """ -$cleanedPlanSummary - -Please see the unit tests generated below. Click 'View Diff' to review the changes in the code editor. - """.trimIndent() // Ensures no extra indentation for the entire message - - val viewDiffMessageId = codeTestChatHelper.addAnswer( - CodeTestChatMessageContent( - message = summary, - type = ChatMessageType.Answer, - buttons = listOf(Button("utg_view_diff", "View Diff", keepCardAfterClick = true, position = "outside", status = "info")), - fileList = listOf(getTestFilePathRelativeToRoot(targetFileInfo)), - projectRootName = project.name, - canBeVoted = true, - codeReference = codeReference - ) - ) - - session.viewDiffMessageId = viewDiffMessageId - codeTestChatHelper.updateUI( - promptInputDisabledState = false, - promptInputPlaceholder = "Specify a function(s) in the current file(optional)", - promptInputProgress = testGenCompletedField, - ) - } else { - codeTestChatHelper.updateAnswer( - CodeTestChatMessageContent( - type = ChatMessageType.Answer, - buttons = listOf(Button("utg_view_diff", "View Diff", keepCardAfterClick = true, position = "outside", status = "info")), - fileList = listOf(getTestFilePathRelativeToRoot(targetFileInfo)), - projectRootName = project.name, - codeReference = codeReference - ), - messageIdOverride = previousIterationContext.buildAndExecuteMessageId - ) - session.viewDiffMessageId = previousIterationContext.buildAndExecuteMessageId - codeTestChatHelper.updateUI( - loadingChat = false, - ) - } - codeTestChatHelper.updateUI( - promptInputDisabledState = true, - promptInputPlaceholder = message("testgen.placeholder.view_diff"), - promptInputProgress = testGenCompletedField, - ) - delay(1000) - } - - codeTestChatHelper.sendUpdatePromptProgress(codeTestChatHelper.getActiveSession().tabId, null) - } - - // Input: test file path relative to project root's parent . - // Output: test file path relative to project root. - // targetFileInfo.testFilePath has a format of /. - // test file path in generatedTestDiffs map has a format of resultArtifacts/. - // both needs to be handled the same way which is remove the first sub-directory - private fun getTestFilePathRelativeToRoot(targetFileInfo: Any?): String { - val pathString = when (targetFileInfo) { - is TargetFileInfo -> targetFileInfo.testFilePath() - else -> generatedTestDiffs.keys.firstOrNull() - } ?: throw RuntimeException("No test file path found") - val path = Paths.get(pathString) - val updatedPath = path.subpath(1, path.nameCount).toString() - return updatedPath - } - - private fun storeGeneratedTestDiffs(byteArray: ByteArray, session: Session) { - try { - val byteArrayInputStream = ByteArrayInputStream(byteArray) - ZipInputStream(byteArrayInputStream).use { zipInputStream -> - var zipEntry = zipInputStream.nextEntry - - while (zipEntry != null) { - if (zipEntry.isDirectory) { - zipInputStream.closeEntry() - zipEntry = zipInputStream.nextEntry - // We are only interested in test file diff in zip entries - continue - } - - val baos = ByteArrayOutputStream() - val buffer = ByteArray(1024) - var len: Int - - while (zipInputStream.read(buffer).also { len = it } > 0) { - baos.write(buffer, 0, len) - } - - val fileContent = baos.toByteArray() - if (fileContent.toString(Charsets.UTF_8).isEmpty()) { - session.isGeneratingTests = false - return - } - val zipEntryPath = Paths.get(zipEntry.name) - - // relative path to project root - val updatedZipEntryPath = zipEntryPath.subpath(1, zipEntryPath.nameCount).toString() - session.generatedTestDiffs[updatedZipEntryPath] = fileContent.toString(Charsets.UTF_8) - - zipInputStream.closeEntry() - zipEntry = zipInputStream.nextEntry - } - } - } catch (e: IOException) { - LOG.debug(e) { "Error reading ZIP entries" } - throw e - } - } - - private suspend fun createUploadUrl( - codeTestChatHelper: CodeTestChatHelper, - previousIterationContext: PreviousUTGIterationContext?, - ): CodeTestResponseContext { - throwIfCancelled(codeTestChatHelper.getActiveSession()) - val file = - if (previousIterationContext == null) { - FileEditorManager.getInstance(project).selectedEditor?.file.also { - codeTestChatHelper.getActiveSession().selectedFile = it - } - } else { - previousIterationContext.selectedFile - } - - val combinedBuildAndExecuteLogFile = combineBuildAndExecuteLogFiles( - previousIterationContext?.buildLogFile, - previousIterationContext?.testLogFile - ) - val codeTestSessionConfig = CodeTestSessionConfig(file, project, combinedBuildAndExecuteLogFile) - codeTestChatHelper.getActiveSession().projectRoot = codeTestSessionConfig.projectRoot.path - - val codeTestSessionContext = CodeTestSessionContext(project, codeTestSessionConfig) - val codeWhispererCodeTestSession = CodeWhispererCodeTestSession(codeTestSessionContext) - return codeWhispererCodeTestSession.run(codeTestChatHelper, previousIterationContext) - } - - private fun startTestGeneration(uploadId: String, targetCode: List, userInput: String): StartTestGenerationResponse = - CodeWhispererClientAdaptor.getInstance(project).startTestGeneration(uploadId, targetCode, userInput) - - private fun getTestGenerationStatus(jobId: String, jobGroupName: String): GetTestGenerationResponse = - CodeWhispererClientAdaptor.getInstance(project).getTestGeneration(jobId, jobGroupName) - - /** - * Returns true if the UTG is in progress. - * This function will return true for a cancelled UTG job which is in cancellation state. - */ - fun isUTGInProgress(): Boolean = isUTGInProgress.get() - - private fun beforeTestGenFlow(session: Session) { - resetTestGenFlowSession(session) - session.isGeneratingTests = true - isUTGInProgress.set(true) - // Show in progress indicator - - ApplicationManager.getApplication().invokeLater { - (FileDocumentManager.getInstance() as FileDocumentManagerImpl).saveAllDocuments(false) - } - } - - private fun resetTestGenFlowSession(session: Session) { - // session.selectedFile doesn't need to be reset since it will remain unchanged - session.conversationState = ConversationState.IN_PROGRESS - session.packageInfoList = emptyList() - session.openedDiffFile = null - session.testFileRelativePathToProjectRoot = "" - session.testFileName = "" - session.openedDiffFile = null - session.generatedTestDiffs.clear() - session.buildAndExecuteTaskContext.apply { - buildExitCode = -1 - testExitCode = -1 - progressStatus = BuildAndExecuteProgressStatus.START_STEP - } - } - - private fun afterTestGenFlow() { - isUTGInProgress.set(false) - } - - /** - * Triggers a unit test generation flow based on current open file. - */ - fun generateTests( - prompt: String, - codeTestChatHelper: CodeTestChatHelper, - previousIterationContext: PreviousUTGIterationContext?, - selectionRange: Range?, - ): Job? { - val shouldStart = performTestGenPreChecks() - val session = codeTestChatHelper.getActiveSession() - if (!shouldStart) { - session.conversationState = ConversationState.IDLE - return null - } - - beforeTestGenFlow(session) - - return cs.launch { - try { - launchTestGenFlow(prompt, codeTestChatHelper, previousIterationContext, selectionRange) - } catch (e: Exception) { - // reset number of unitTestGenerated to null - session.numberOfUnitTestCasesGenerated = null - // Add an answer for displaying error message - val errorMessage = when { - e is CodeTestException && - e.message?.startsWith("CreateTestJobError: Maximum") == true -> - message("testgen.error.maximum_generations_reach") - - e is CodeTestException -> e.uiMessage - e is JsonParseException -> message("testgen.error.generic_technical_error_message") - else -> message("testgen.error.generic_error_message") - } - val buttonList = mutableListOf