diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt deleted file mode 100644 index c9fe08efc64..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt +++ /dev/null @@ -1,113 +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.amazonqFeatureDev - -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.project.Project -import kotlinx.coroutines.launch -import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope -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.amazonqDoc.auth.isDocAvailable -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.auth.isFeatureDevAvailable -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.controller.FeatureDevController -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.AuthenticationUpdateMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.IncomingFeatureDevMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage -import software.aws.toolkits.jetbrains.services.codemodernizer.utils.isCodeTransformAvailable - -class FeatureDevApp : AmazonQApp { - - private val scope = disposableCoroutineScope(this) - - override val tabTypes = listOf("featuredev") - - override fun init(context: AmazonQAppInitContext) { - val chatSessionStorage = ChatSessionStorage() - // Create FeatureDev controller - val inboundAppMessagesHandler = - FeatureDevController(context, chatSessionStorage) - - context.messageTypeRegistry.register( - "chat-prompt" to IncomingFeatureDevMessage.ChatPrompt::class, - "new-tab-was-created" to IncomingFeatureDevMessage.NewTabCreated::class, - "tab-was-removed" to IncomingFeatureDevMessage.TabRemoved::class, - "auth-follow-up-was-clicked" to IncomingFeatureDevMessage.AuthFollowUpWasClicked::class, - "follow-up-was-clicked" to IncomingFeatureDevMessage.FollowupClicked::class, - "chat-item-voted" to IncomingFeatureDevMessage.ChatItemVotedMessage::class, - "chat-item-feedback" to IncomingFeatureDevMessage.ChatItemFeedbackMessage::class, - "response-body-link-click" to IncomingFeatureDevMessage.ClickedLink::class, - "insert_code_at_cursor_position" to IncomingFeatureDevMessage.InsertCodeAtCursorPosition::class, - "open-diff" to IncomingFeatureDevMessage.OpenDiff::class, - "file-click" to IncomingFeatureDevMessage.FileClicked::class, - "stop-response" to IncomingFeatureDevMessage.StopResponse::class, - "store-code-result-message-id" to IncomingFeatureDevMessage.StoreMessageIdMessage::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 IncomingFeatureDevMessage.ChatPrompt -> inboundAppMessagesHandler.processPromptChatMessage(message) - is IncomingFeatureDevMessage.NewTabCreated -> inboundAppMessagesHandler.processNewTabCreatedMessage(message) - is IncomingFeatureDevMessage.TabRemoved -> inboundAppMessagesHandler.processTabRemovedMessage(message) - is IncomingFeatureDevMessage.AuthFollowUpWasClicked -> inboundAppMessagesHandler.processAuthFollowUpClick(message) - is IncomingFeatureDevMessage.FollowupClicked -> inboundAppMessagesHandler.processFollowupClickedMessage(message) - is IncomingFeatureDevMessage.ChatItemVotedMessage -> inboundAppMessagesHandler.processChatItemVotedMessage(message) - is IncomingFeatureDevMessage.ChatItemFeedbackMessage -> inboundAppMessagesHandler.processChatItemFeedbackMessage(message) - is IncomingFeatureDevMessage.ClickedLink -> inboundAppMessagesHandler.processLinkClick(message) - is IncomingFeatureDevMessage.InsertCodeAtCursorPosition -> inboundAppMessagesHandler.processInsertCodeAtCursorPosition(message) - is IncomingFeatureDevMessage.OpenDiff -> inboundAppMessagesHandler.processOpenDiff(message) - is IncomingFeatureDevMessage.FileClicked -> inboundAppMessagesHandler.processFileClicked(message) - is IncomingFeatureDevMessage.StopResponse -> inboundAppMessagesHandler.processStopMessage(message) - is IncomingFeatureDevMessage.StoreMessageIdMessage -> inboundAppMessagesHandler.processStoreCodeResultMessageId(message) - } - } - - override fun dispose() { - // nothing to do - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevAppFactory.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevAppFactory.kt deleted file mode 100644 index 9aa279c67b8..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevAppFactory.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.amazonqFeatureDev - -import com.intellij.openapi.project.Project -import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppFactory - -class FeatureDevAppFactory : AmazonQAppFactory { - override fun createApp(project: Project) = FeatureDevApp() -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevConstants.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevConstants.kt deleted file mode 100644 index dc15a5c870c..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevConstants.kt +++ /dev/null @@ -1,67 +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.amazonqFeatureDev - -const val FEATURE_EVALUATION_PRODUCT_NAME = "FeatureDev" - -const val FEATURE_NAME = "Amazon Q Developer Agent for software development" - -@Suppress("MaxLineLength") -const val GENERATE_DEV_FILE_PROMPT = "generate a devfile in my repository. Note that you should only use devfile version 2.0.0 and the only supported commands are install, build and test (are all optional). so you may have to bundle some commands together using '&&'. also you can use \"public.ecr.aws/aws-mde/universal-image:latest\" as universal image if you aren’t sure which image to use. here is an example for a node repository (but don't assume it's always a node project. look at the existing repository structure before generating the devfile): schemaVersion: 2.0.0 components: - name: dev container: image: public.ecr.aws/aws-mde/universal-image:latest commands: - id: install exec: component: dev commandLine: \"npm install\" - id: build exec: component: dev commandLine: \"npm run build\" - id: test exec: component: dev commandLine: \"npm run test\"" - -// Max number of times a user can attempt to retry a code generation request if it fails -const val CODE_GENERATION_RETRY_LIMIT = 3 - -// The default retry limit used when the session could not be found -const val DEFAULT_RETRY_LIMIT = 0 - -// Max allowed size for a repository in bytes -const val MAX_PROJECT_SIZE_BYTES: Long = 200 * 1024 * 1024 - -val CLIENT_ERROR_MESSAGES = setOf( - "StartTaskAssistCodeGeneration reached for this month.", - "The folder you chose did not contain any source files in a supported language. Choose another folder and try again.", - "reached the quota for number of iterations on code generation." -) - -enum class ModifySourceFolderErrorReason( - private val reasonText: String, -) { - ClosedBeforeSelection("ClosedBeforeSelection"), - NotInWorkspaceFolder("NotInWorkspaceFolder"), - ; - - override fun toString(): String = reasonText -} - -enum class FeatureDevOperation(private val operationName: String) { - StartTaskAssistCodeGeneration("StartTaskAssistCodeGenerator"), - CreateConversation("CreateConversation"), - CreateUploadUrl("CreateUploadUrl"), - GenerateCode("GenerateCode"), - GetTaskAssistCodeGeneration("GetTaskAssistCodeGenerator"), - ExportTaskAssistArchiveResult("ExportTaskAssistArchiveResult"), - UploadToS3("UploadToS3"), - ; - - override fun toString(): String = operationName -} - -enum class MetricDataOperationName(private val operationName: String) { - StartCodeGeneration("StartCodeGeneration"), - EndCodeGeneration("EndCodeGeneration"), - ; - - override fun toString(): String = operationName -} - -enum class MetricDataResult(private val resultName: String) { - Success("Success"), - Fault("Fault"), - Error("Error"), - LlmFailure("LLMFailure"), - ; - - override fun toString(): String = resultName -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevExceptions.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevExceptions.kt deleted file mode 100644 index add731461c0..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevExceptions.kt +++ /dev/null @@ -1,105 +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.amazonqFeatureDev - -import software.aws.toolkits.jetbrains.services.amazonq.project.RepoSizeError -import software.aws.toolkits.resources.message - -/** - * FeatureDevException models failures from feature dev operations. - * - * - Each failure is annotated based on className, operation, and a short desc. Use the `reason()` and `reasonDesc()` members for instrumentation. - * - To throw an exception without modeling, throw FeatureDevException directly. - */ -open class FeatureDevException(override val message: String?, val operation: String, val desc: String?, override val cause: Throwable? = null) : - RuntimeException() { - fun reason(): String = this.javaClass.simpleName - - fun reasonDesc(): String = - when { - !desc.isNullOrEmpty() -> "$operation | Description: $desc" - !message.isNullOrEmpty() -> "$operation | Description: $message" - else -> operation - } -} - -/** - * Exceptions extending this class are considered "errors" in service metrics. - */ -open class ClientException(message: String, operation: String, desc: String?, cause: Throwable? = null) : - FeatureDevException(message, operation, desc, cause) - -/** - * Errors extending this class are considered "faults" in service metrics. - */ -open class ServiceException(message: String, operation: String, desc: String?, cause: Throwable? = null) : - FeatureDevException(message, operation, desc, cause) - -/** - * Errors extending this class are considered "LLM failures" in service metrics. - */ -open class LlmException(message: String, operation: String, desc: String?, cause: Throwable? = null) : - FeatureDevException(message, operation, desc, cause) - -object ApiException { - fun of(statusCode: Int, message: String, operation: String, desc: String?, cause: Throwable? = null): FeatureDevException = - when (statusCode in 400..499) { - true -> ClientException(message, operation, desc, cause) - false -> ServiceException(message, operation, desc, cause) - } -} - -class NoChangeRequiredException(operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message("amazonqFeatureDev.exception.no_change_required_exception"), operation, desc, cause) - -class EmptyPatchException(operation: String, desc: String?, cause: Throwable? = null) : - LlmException(message("amazonqFeatureDev.exception.guardrails"), operation, desc, cause) - -class ContentLengthException( - override val message: String = message("amazonqFeatureDev.content_length.error_text"), - operation: String, - desc: String?, - cause: Throwable? = null, -) : - RepoSizeError, ClientException(message, operation, desc, cause) - -class ZipFileCorruptedException(operation: String, desc: String?, cause: Throwable? = null) : - ServiceException("The zip file is corrupted", operation, desc, cause) - -class UploadURLExpired(operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message("amazonqFeatureDev.exception.upload_url_expiry"), operation, desc, cause) - -class CodeIterationLimitException(operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message("amazonqFeatureDev.code_generation.iteration_limit.error_text"), operation, desc, cause) - -class MonthlyConversationLimitError(message: String, operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message, operation, desc, cause) - -class GuardrailsException(operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message("amazonqFeatureDev.exception.guardrails"), operation, desc, cause) - -class PromptRefusalException(operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message("amazonqFeatureDev.exception.prompt_refusal"), operation, desc, cause) - -class FileCreationFailedException(operation: String, desc: String?, cause: Throwable? = null) : - ServiceException(message("amazonqFeatureDev.exception.failed_generation"), operation, desc, cause) - -class ThrottlingException(operation: String, desc: String?, cause: Throwable? = null) : - ClientException(message("amazonqFeatureDev.exception.throttling"), operation, desc, cause) - -class ExportParseException(operation: String, desc: String?, cause: Throwable? = null) : - ServiceException(message("amazonqFeatureDev.exception.export_parsing_error"), operation, desc, cause) - -class CodeGenerationException(operation: String, desc: String?, cause: Throwable? = null) : - ServiceException(message("amazonqFeatureDev.code_generation.failed_generation"), operation, desc, cause) - -class UploadCodeException(operation: String, desc: String?, cause: Throwable? = null) : - ServiceException(message("amazonqFeatureDev.exception.upload_code"), operation, desc, cause) - -class ConversationIdNotFoundException(operation: String, desc: String?, cause: Throwable? = null) : - ServiceException(message("amazonqFeatureDev.exception.conversation_not_found"), operation, desc, cause) - -val denyListedErrors = arrayOf("Deserialization error", "Inaccessible host", "UnknownHost") -fun createUserFacingErrorMessage(message: String?): String? = - if (message != null && denyListedErrors.any { message.contains(it) }) "$FEATURE_NAME API request failed" else message diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/InboundAppMessagesHandler.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/InboundAppMessagesHandler.kt deleted file mode 100644 index 2c4d8cd0e16..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/InboundAppMessagesHandler.kt +++ /dev/null @@ -1,23 +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.amazonqFeatureDev - -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.IncomingFeatureDevMessage - -interface InboundAppMessagesHandler { - - suspend fun processPromptChatMessage(message: IncomingFeatureDevMessage.ChatPrompt) - suspend fun processNewTabCreatedMessage(message: IncomingFeatureDevMessage.NewTabCreated) - suspend fun processTabRemovedMessage(message: IncomingFeatureDevMessage.TabRemoved) - suspend fun processAuthFollowUpClick(message: IncomingFeatureDevMessage.AuthFollowUpWasClicked) - suspend fun processFollowupClickedMessage(message: IncomingFeatureDevMessage.FollowupClicked) - suspend fun processChatItemVotedMessage(message: IncomingFeatureDevMessage.ChatItemVotedMessage) - suspend fun processChatItemFeedbackMessage(message: IncomingFeatureDevMessage.ChatItemFeedbackMessage) - suspend fun processLinkClick(message: IncomingFeatureDevMessage.ClickedLink) - suspend fun processInsertCodeAtCursorPosition(message: IncomingFeatureDevMessage.InsertCodeAtCursorPosition) - suspend fun processOpenDiff(message: IncomingFeatureDevMessage.OpenDiff) - suspend fun processFileClicked(message: IncomingFeatureDevMessage.FileClicked) - suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse) - suspend fun processStoreCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt deleted file mode 100644 index 7b25fb3b00e..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt +++ /dev/null @@ -1,244 +0,0 @@ -// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.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.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.services.amazonq.clients.AmazonQStreamingClient -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.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 FeatureDevClient( - private val project: Project, -) { - fun getTelemetryOptOutPreference() = - if (AwsSettings.getInstance().isTelemetryEnabled) { - OptOutPreference.OPTIN - } else { - OptOutPreference.OPTOUT - } - - private val featureDevUserContext = - 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() - } - - private fun bearerClient() = QRegionProfileManager.getInstance().getQClient(project) - - private val amazonQStreamingClient - get() = AmazonQStreamingClient.getInstance(project) - - fun sendFeatureDevTelemetryEvent(conversationId: String): SendTelemetryEventResponse = - bearerClient().sendTelemetryEvent { requestBuilder -> - requestBuilder.telemetryEvent { telemetryEventBuilder -> - telemetryEventBuilder.featureDevEvent { - it.conversationId(conversationId) - } - } - requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(featureDevUserContext) - requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun sendFeatureDevMetricData(operationName: String, result: String): SendTelemetryEventResponse = - bearerClient().sendTelemetryEvent { requestBuilder -> - requestBuilder.telemetryEvent { telemetryEventBuilder -> - telemetryEventBuilder.metricData { - it - .metricName("Operation") - .metricValue(1.0) - .timestamp(Instant.now()) - .product("FeatureDev") - .dimensions( - listOf( - Dimension.builder() - .name("operationName") - .value(operationName) - .build(), - Dimension.builder() - .name("result") - .value(result) - .build() - ) - ) - } - } - requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(featureDevUserContext) - requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun sendFeatureDevCodeGenerationEvent( - conversationId: String, - linesOfCodeGenerated: Int, - charactersOfCodeGenerated: Int, - ): SendTelemetryEventResponse = - bearerClient().sendTelemetryEvent { requestBuilder -> - requestBuilder.telemetryEvent { telemetryEventBuilder -> - telemetryEventBuilder.featureDevCodeGenerationEvent { - it - .conversationId(conversationId) - .linesOfCodeGenerated(linesOfCodeGenerated) - .charactersOfCodeGenerated(charactersOfCodeGenerated) - } - } - requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(featureDevUserContext) - requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun sendFeatureDevCodeAcceptanceEvent( - conversationId: String, - linesOfCodeAccepted: Int, - charactersOfCodeAccepted: Int, - ): SendTelemetryEventResponse = - bearerClient().sendTelemetryEvent { requestBuilder -> - requestBuilder.telemetryEvent { telemetryEventBuilder -> - telemetryEventBuilder.featureDevCodeAcceptanceEvent { - it - .conversationId(conversationId) - .linesOfCodeAccepted(linesOfCodeAccepted) - .charactersOfCodeAccepted(charactersOfCodeAccepted) - } - } - requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(featureDevUserContext) - requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun createTaskAssistConversation(): CreateTaskAssistConversationResponse = - bearerClient().createTaskAssistConversation( - CreateTaskAssistConversationRequest.builder() - .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - .build(), - ) - - fun createTaskAssistUploadUrl( - conversationId: String, - contentChecksumSha256: String, - contentLength: Long, - uploadId: String, - ): CreateUploadUrlResponse = - bearerClient().createUploadUrl { - it - .contentChecksumType(ContentChecksumType.SHA_256) - .uploadId(uploadId) - .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, - codeGenerationId: String?, - currentCodeGenerationId: String?, - ): 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") } // This parameter is omitted by featureDev but required in the request - .uploadId(uploadId) - }.codeGenerationId(codeGenerationId.toString()) - .currentCodeGenerationId(currentCodeGenerationId) - .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - fun getTaskAssistCodeGeneration( - conversationId: String, - codeGenerationId: String, - ): GetTaskAssistCodeGenerationResponse = - bearerClient() - .getTaskAssistCodeGeneration { - it - .conversationId(conversationId) - .codeGenerationId(codeGenerationId) - .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) - } - - 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/services/amazonqFeatureDev/controller/FeatureDevController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt deleted file mode 100644 index 89171dab382..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt +++ /dev/null @@ -1,921 +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.amazonqFeatureDev.controller - -import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper -import com.intellij.diff.DiffContentFactory -import com.intellij.diff.chains.SimpleDiffRequestChain -import com.intellij.diff.contents.EmptyContent -import com.intellij.diff.editor.ChainDiffVirtualFile -import com.intellij.diff.editor.DiffEditorTabFilesManager -import com.intellij.diff.requests.SimpleDiffRequest -import com.intellij.ide.BrowserUtil -import com.intellij.openapi.application.runInEdt -import com.intellij.openapi.command.WriteCommandAction -import com.intellij.openapi.editor.Caret -import com.intellij.openapi.editor.Editor -import com.intellij.openapi.fileEditor.FileEditorManager -import com.intellij.openapi.project.Project -import com.intellij.openapi.vfs.VfsUtil -import com.intellij.openapi.wm.ToolWindowManager -import kotlinx.coroutines.delay -import kotlinx.coroutines.withContext -import software.amazon.awssdk.services.toolkittelemetry.model.Sentiment -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.core.utils.warn -import software.aws.toolkits.jetbrains.common.util.selectFolder -import software.aws.toolkits.jetbrains.core.coroutines.EDT -import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext -import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener -import software.aws.toolkits.jetbrains.services.amazonq.project.RepoSizeError -import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeIterationLimitException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.DEFAULT_RETRY_LIMIT -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.GENERATE_DEV_FILE_PROMPT -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.InboundAppMessagesHandler -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ModifySourceFolderErrorReason -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MonthlyConversationLimitError -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.NoChangeRequiredException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.UploadURLExpired -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ZipFileCorruptedException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.createUserFacingErrorMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.denyListedErrors -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FeatureDevMessageType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpIcons -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpStatusType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpTypes -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.IncomingFeatureDevMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.initialExamples -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswer -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAuthNeededException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAuthenticationInProgressMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendChatInputEnabledMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendError -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendMonthlyLimitError -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendSystemPrompt -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateFileComponent -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions -import software.aws.toolkits.jetbrains.services.codewhisperer.util.content -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.FeedbackComment -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl -import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService -import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings -import software.aws.toolkits.jetbrains.ui.feedback.FeatureDevFeedbackDialog -import software.aws.toolkits.jetbrains.utils.notifyError -import software.aws.toolkits.resources.message -import software.aws.toolkits.telemetry.AmazonqTelemetry -import software.aws.toolkits.telemetry.Result -import software.aws.toolkits.telemetry.UiTelemetry -import java.util.UUID - -class FeatureDevController( - private val context: AmazonQAppInitContext, - private val chatSessionStorage: ChatSessionStorage, - private val authController: AuthController = AuthController(), -) : InboundAppMessagesHandler { - - init { - context.project.messageBus.connect().subscribe( - QRegionProfileSelectedListener.TOPIC, - object : QRegionProfileSelectedListener { - override fun onProfileSelected(project: Project, profile: QRegionProfile?) { - chatSessionStorage.deleteAllSessions() - } - } - ) - } - - val messenger = context.messagesFromAppToUi - val toolWindow = ToolWindowManager.getInstance(context.project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID) - - private val diffVirtualFiles = mutableMapOf() - - override suspend fun processPromptChatMessage(message: IncomingFeatureDevMessage.ChatPrompt) { - handleChat( - tabId = message.tabId, - message = message.chatMessage - ) - } - - override suspend fun processStoreCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) { - storeCodeResultMessageId(message) - } - - override suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse) { - handleStopMessage(message) - } - - override suspend fun processNewTabCreatedMessage(message: IncomingFeatureDevMessage.NewTabCreated) { - newTabOpened(message.tabId) - } - - override suspend fun processTabRemovedMessage(message: IncomingFeatureDevMessage.TabRemoved) { - chatSessionStorage.deleteSession(message.tabId) - } - - override suspend fun processAuthFollowUpClick(message: IncomingFeatureDevMessage.AuthFollowUpWasClicked) { - authController.handleAuth(context.project, message.authType) - messenger.sendAuthenticationInProgressMessage(message.tabId) // show user that authentication is in progress - messenger.sendChatInputEnabledMessage(message.tabId, enabled = false) // disable the input field while authentication is in progress - } - - override suspend fun processFollowupClickedMessage(message: IncomingFeatureDevMessage.FollowupClicked) { - when (message.followUp.type) { - FollowUpTypes.RETRY -> retryRequests(message.tabId) - FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER -> modifyDefaultSourceFolder(message.tabId) - FollowUpTypes.DEV_EXAMPLES -> messenger.initialExamples(message.tabId) - FollowUpTypes.SEND_FEEDBACK -> sendFeedback() - FollowUpTypes.INSERT_CODE -> insertCode(message.tabId) - FollowUpTypes.PROVIDE_FEEDBACK_AND_REGENERATE_CODE -> provideFeedbackAndRegenerateCode(message.tabId) - FollowUpTypes.NEW_TASK -> newTask(message.tabId) - FollowUpTypes.CLOSE_SESSION -> closeSession(message.tabId) - FollowUpTypes.ACCEPT_AUTO_BUILD -> handleDevCommandUserSetting(message.tabId, true) - FollowUpTypes.DENY_AUTO_BUILD -> handleDevCommandUserSetting(message.tabId, false) - FollowUpTypes.GENERATE_DEV_FILE -> { - messenger.sendAnswer( - tabId = message.tabId, - messageType = FeatureDevMessageType.SystemPrompt, - message = message("amazonqFeatureDev.follow_up.generate_dev_file") - ) - newTask(tabId = message.tabId, prefilledPrompt = GENERATE_DEV_FILE_PROMPT) - } - } - } - - override suspend fun processChatItemVotedMessage(message: IncomingFeatureDevMessage.ChatItemVotedMessage) { - logger.debug { "$FEATURE_NAME: Processing ChatItemVotedMessage: $message" } - this.disablePreviousFileList(message.tabId) - - val session = chatSessionStorage.getSession(message.tabId, context.project) - when (message.vote) { - "upvote" -> { - AmazonqTelemetry.codeGenerationThumbsUp( - amazonqConversationId = session.conversationId, - credentialStartUrl = getStartUrl(project = context.project) - ) - } - "downvote" -> { - AmazonqTelemetry.codeGenerationThumbsDown( - amazonqConversationId = session.conversationId, - credentialStartUrl = getStartUrl(project = context.project) - ) - } - } - } - - override suspend fun processChatItemFeedbackMessage(message: IncomingFeatureDevMessage.ChatItemFeedbackMessage) { - logger.debug { "$FEATURE_NAME: Processing ChatItemFeedbackMessage: ${message.comment}" } - - val session = getSessionInfo(message.tabId) - - val comment = FeedbackComment( - conversationId = session.conversationId, - userComment = message.comment.orEmpty(), - reason = message.selectedOption, - messageId = message.messageId, - type = "featuredev-chat-answer-feedback" - ) - - try { - TelemetryService.getInstance().sendFeedback( - sentiment = Sentiment.NEGATIVE, - comment = objectMapper.writeValueAsString(comment), - ) - logger.info { "$FEATURE_NAME answer feedback sent: \"Negative\"" } - } catch (e: Throwable) { - e.notifyError(message("feedback.submit_failed", e)) - logger.warn(e) { "Failed to submit feedback" } - return - } - } - - override suspend fun processLinkClick(message: IncomingFeatureDevMessage.ClickedLink) { - BrowserUtil.browse(message.link) - } - - override suspend fun processInsertCodeAtCursorPosition(message: IncomingFeatureDevMessage.InsertCodeAtCursorPosition) { - logger.debug { "$FEATURE_NAME: Processing InsertCodeAtCursorPosition: $message" } - - withContext(EDT) { - val editor: Editor = FileEditorManager.getInstance(context.project).selectedTextEditorWithRemotes.firstOrNull() ?: return@withContext - - val caret: Caret = editor.caretModel.primaryCaret - val offset: Int = caret.offset - - WriteCommandAction.runWriteCommandAction(context.project) { - if (caret.hasSelection()) { - editor.document.deleteString(caret.selectionStart, caret.selectionEnd) - } - editor.document.insertString(offset, message.code) - } - } - } - - private fun putDiff(filePath: String, request: SimpleDiffRequest) { - // Close any existing diff and open a new diff, as the diff virtual file does not appear to allow replacing content directly: - val existingDiff = diffVirtualFiles[filePath] - if (existingDiff != null) { - FileEditorManager.getInstance(context.project).closeFile(existingDiff) - } - - val newDiff = ChainDiffVirtualFile(SimpleDiffRequestChain(request), filePath) - DiffEditorTabFilesManager.getInstance(context.project).showDiffFile(newDiff, true) - diffVirtualFiles[filePath] = newDiff - } - - override suspend fun processOpenDiff(message: IncomingFeatureDevMessage.OpenDiff) { - val session = getSessionInfo(message.tabId) - - AmazonqTelemetry.isReviewedChanges( - amazonqConversationId = session.conversationId, - enabled = true, - credentialStartUrl = getStartUrl(project = context.project) - ) - - val project = context.project - val sessionState = session.sessionState - - when (sessionState) { - is PrepareCodeGenerationState -> { - runInEdt { - val existingFile = VfsUtil.findRelativeFile(message.filePath, session.context.addressableRoot) - - val leftDiffContent = if (existingFile == null) { - EmptyContent() - } else { - DiffContentFactory.getInstance().create(project, existingFile) - } - - val newFileContent = sessionState.filePaths.find { it.zipFilePath == message.filePath }?.fileContent - - val rightDiffContent = if (message.deleted || newFileContent == null) { - EmptyContent() - } else { - DiffContentFactory.getInstance().create(newFileContent) - } - - putDiff(message.filePath, SimpleDiffRequest(message.filePath, leftDiffContent, rightDiffContent, null, null)) - } - } - else -> { - logger.error { "$FEATURE_NAME: OpenDiff event is received for a conversation that has ${session.sessionState.phase} phase" } - messenger.sendError( - tabId = message.tabId, - errMessage = message("amazonqFeatureDev.exception.open_diff_failed"), - retries = 0, - conversationId = session.conversationIdUnsafe - ) - } - } - } - - override suspend fun processFileClicked(message: IncomingFeatureDevMessage.FileClicked) { - val fileToUpdate = message.filePath - val session = getSessionInfo(message.tabId) - val messageId = message.messageId - val action = message.actionName - - var filePaths: List = emptyList() - var deletedFiles: List = emptyList() - var references: List = emptyList() - when (val state = session.sessionState) { - is PrepareCodeGenerationState -> { - filePaths = state.filePaths - deletedFiles = state.deletedFiles - references = state.references - } - } - - fun insertAction(): InsertAction = - if (filePaths.all { it.changeApplied } && deletedFiles.all { it.changeApplied }) { - InsertAction.AUTO_CONTINUE - } else if (filePaths.all { it.changeApplied || it.rejected } && deletedFiles.all { it.changeApplied || it.rejected }) { - InsertAction.CONTINUE - } else if (filePaths.any { it.changeApplied || it.rejected } || deletedFiles.any { it.changeApplied || it.rejected }) { - InsertAction.REMAINING - } else { - InsertAction.ALL - } - - val prevInsertAction = insertAction() - - if (action == "accept-change") { - session.insertChanges( - filePaths = filePaths.filter { it.zipFilePath == fileToUpdate }, - deletedFiles = deletedFiles.filter { it.zipFilePath == fileToUpdate }, - references = references, // Add all references (not attributed per-file) - ) - - AmazonqTelemetry.isAcceptedCodeChanges( - amazonqNumberOfFilesAccepted = 1.0, - amazonqConversationId = session.conversationId, - enabled = true, - credentialStartUrl = getStartUrl(project = context.project) - ) - } else { - // Mark the file as rejected or not depending on the previous state - filePaths.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected } - deletedFiles.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected } - } - - messenger.updateFileComponent(message.tabId, filePaths, deletedFiles, messageId) - - // Then, if the accepted file is not a deletion, open a diff to show the changes are applied: - if (action == "accept-change" && deletedFiles.none { it.zipFilePath == fileToUpdate }) { - var pollAttempt = 0 - val pollDelayMs = 10L - while (pollAttempt < 5) { - val file = VfsUtil.findRelativeFile(message.filePath, session.context.addressableRoot) - // Wait for the file to be created and/or updated to the new content: - if (file != null && file.content() == filePaths.find { it.zipFilePath == fileToUpdate }?.fileContent) { - // Open a diff, showing the changes have been applied and the file now has identical left/right state: - this.processOpenDiff(IncomingFeatureDevMessage.OpenDiff(message.tabId, fileToUpdate, false)) - break - } else { - pollAttempt++ - delay(pollDelayMs) - } - } - } - - val nextInsertAction = insertAction() - if (nextInsertAction == InsertAction.AUTO_CONTINUE) { - // Insert remaining changes (noop, as there are none), and advance to the next prompt: - insertCode(message.tabId) - } else if (nextInsertAction != prevInsertAction) { - // Update the action displayed to the customer based on the current state: - messenger.sendSystemPrompt(message.tabId, getFollowUpOptions(session.sessionState.phase, nextInsertAction)) - } - } - - private suspend fun newTabOpened(tabId: String) { - var session: Session? = null - try { - session = getSessionInfo(tabId) - logger.debug { "$FEATURE_NAME: Session created with id: ${session.tabID}" } - - val credentialState = authController.getAuthNeededStates(context.project).amazonQ - if (credentialState != null) { - messenger.sendAuthNeededException( - tabId = tabId, - triggerId = UUID.randomUUID().toString(), - credentialState = credentialState, - ) - session.isAuthenticating = true - return - } - } catch (err: Exception) { - val message = createUserFacingErrorMessage(err.message) - messenger.sendError( - tabId = tabId, - errMessage = message ?: message("amazonqFeatureDev.exception.request_failed"), - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe - ) - } - } - - private suspend fun handleStopMessage(message: IncomingFeatureDevMessage.StopResponse) { - val session: Session? - UiTelemetry.click(null as Project?, "amazonq_stopCodeGeneration") - messenger.sendAnswer( - tabId = message.tabId, - message("amazonqFeatureDev.code_generation.stopping_code_generation"), - messageType = FeatureDevMessageType.Answer, - canBeVoted = false - ) - messenger.sendUpdatePlaceholder( - tabId = message.tabId, - newPlaceholder = message("amazonqFeatureDev.code_generation.stopping_code_generation") - ) - messenger.sendChatInputEnabledMessage(tabId = message.tabId, enabled = false) - session = getSessionInfo(message.tabId) - - if (session.sessionState.token?.token !== null) { - session.sessionState.token?.cancel() - } - } - - suspend fun insertCode(tabId: String) { - var session: Session? = null - try { - session = getSessionInfo(tabId) - - var filePaths: List = emptyList() - var deletedFiles: List = emptyList() - var references: List = emptyList() - - when (val state = session.sessionState) { - is PrepareCodeGenerationState -> { - filePaths = state.filePaths - deletedFiles = state.deletedFiles - references = state.references - } - } - - val rejectedFilesCount = filePaths.count { it.rejected } + deletedFiles.count { it.rejected } - val acceptedFilesCount = filePaths.count { it.changeApplied } + filePaths.count { it.changeApplied } - val remainingFilesCount = filePaths.count() + deletedFiles.count() - acceptedFilesCount - rejectedFilesCount - - AmazonqTelemetry.isAcceptedCodeChanges( - amazonqNumberOfFilesAccepted = remainingFilesCount.toDouble(), - amazonqConversationId = session.conversationId, - enabled = true, - credentialStartUrl = getStartUrl(project = context.project) - ) - - session.insertChanges( - filePaths = filePaths, - deletedFiles = deletedFiles, - references = references - ) - session.updateFilesPaths( - filePaths = filePaths, - deletedFiles = deletedFiles, - messenger - ) - - messenger.sendAnswer( - tabId = tabId, - message = message("amazonqFeatureDev.code_generation.updated_code"), - messageType = FeatureDevMessageType.Answer, - canBeVoted = true - ) - - val followUps = mutableListOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.new_task"), - type = FollowUpTypes.NEW_TASK, - status = FollowUpStatusType.Info - ), - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.close_session"), - type = FollowUpTypes.CLOSE_SESSION, - status = FollowUpStatusType.Info - ), - ) - - if (!session.context.hasDevFile()) { - followUps.add( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.generate_dev_file"), - type = FollowUpTypes.GENERATE_DEV_FILE, - status = FollowUpStatusType.Info - ) - ) - - messenger.sendAnswer( - tabId = tabId, - message = message("amazonqFeatureDev.chat_message.generate_dev_file"), - messageType = FeatureDevMessageType.Answer - ) - } - - messenger.sendSystemPrompt( - tabId = tabId, - followUp = followUps - ) - - messenger.sendUpdatePlaceholder( - tabId = tabId, - newPlaceholder = message("amazonqFeatureDev.placeholder.additional_improvements") - ) - } catch (err: Exception) { - val message = createUserFacingErrorMessage("Failed to insert code changes: ${err.message}") - messenger.sendError( - tabId = tabId, - errMessage = message ?: message("amazonqFeatureDev.exception.insert_code_failed"), - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe - ) - } - } - - private suspend fun newTask(tabId: String, isException: Boolean? = false, prefilledPrompt: String? = null) { - val session = getSessionInfo(tabId) - val sessionLatency = System.currentTimeMillis() - session.sessionStartTime - - AmazonqTelemetry.endChat( - amazonqConversationId = session.conversationId, - amazonqEndOfTheConversationLatency = sessionLatency.toDouble(), - credentialStartUrl = getStartUrl(project = context.project) - ) - chatSessionStorage.deleteSession(tabId) - - newTabOpened(tabId) - - if (prefilledPrompt != null && isException != null && !isException) { - handleChat(tabId = tabId, message = prefilledPrompt) - } else { - if (isException != null && !isException) { - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.chat_message.ask_for_new_task") - ) - } - messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.new_plan")) - messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true) - } - } - - private suspend fun handleDevCommandUserSetting(tabId: String, value: Boolean) { - val session = getSessionInfo(tabId) - CodeWhispererSettings.getInstance().toggleAutoBuildFeature(session.context.workspaceRoot.path, value) - messenger.sendAnswer( - tabId = tabId, - message = message("amazonqFeatureDev.chat_message.setting_updated"), - messageType = FeatureDevMessageType.Answer, - ) - this.retryRequests(tabId) - } - - private suspend fun closeSession(tabId: String) { - this.disablePreviousFileList(tabId) - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.chat_message.closed_session"), - canBeVoted = true - ) - - messenger.sendUpdatePlaceholder( - tabId = tabId, - newPlaceholder = message("amazonqFeatureDev.placeholder.closed_session") - ) - - messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) - - val session = getSessionInfo(tabId) - val sessionLatency = System.currentTimeMillis() - session.sessionStartTime - AmazonqTelemetry.endChat( - amazonqConversationId = session.conversationId, - amazonqEndOfTheConversationLatency = sessionLatency.toDouble(), - credentialStartUrl = getStartUrl(project = context.project) - ) - } - - private suspend fun provideFeedbackAndRegenerateCode(tabId: String) { - val session = getSessionInfo(tabId) - - AmazonqTelemetry.isProvideFeedbackForCodeGen( - amazonqConversationId = session.conversationId, - enabled = true, - credentialStartUrl = getStartUrl(project = context.project) - ) - - // Unblock the message button - messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) - - messenger.sendAnswer( - tabId = tabId, - message = message("amazonqFeatureDev.code_generation.provide_code_feedback"), - messageType = FeatureDevMessageType.Answer, - canBeVoted = true - ) - messenger.sendUpdatePlaceholder(tabId, message("amazonqFeatureDev.placeholder.provide_code_feedback")) - } - - private suspend fun processErrorChatMessage(err: Exception, message: String, session: Session?, tabId: String) { - logger.warn(err) { "Encountered ${err.message} for tabId: $tabId" } - when (err) { - is RepoSizeError -> { - messenger.sendError( - tabId = tabId, - errMessage = err.message, - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe - ) - messenger.sendSystemPrompt( - tabId = tabId, - followUp = listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"), - type = FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, - status = FollowUpStatusType.Info, - ) - ), - ) - } - is NoChangeRequiredException -> { - val isException = true - messenger.sendAnswer( - tabId = tabId, - message = err.message, - messageType = FeatureDevMessageType.Answer, - canBeVoted = true - ) - return this.newTask(message, isException) - } - is ZipFileCorruptedException -> { - messenger.sendError( - tabId = tabId, - errMessage = err.message, - retries = 0, - conversationId = session?.conversationIdUnsafe - ) - } - is MonthlyConversationLimitError -> { - messenger.sendMonthlyLimitError(tabId = tabId) - messenger.sendChatInputEnabledMessage(tabId, enabled = false) - } - is UploadURLExpired -> messenger.sendAnswer( - tabId = tabId, - message = err.message, - messageType = FeatureDevMessageType.Answer, - canBeVoted = true - ) - is CodeIterationLimitException -> { - messenger.sendError( - tabId = tabId, - errMessage = err.message, - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe - ) - messenger.sendSystemPrompt( - tabId = tabId, - followUp = listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.insert_all_code"), - type = FollowUpTypes.INSERT_CODE, - icon = FollowUpIcons.Ok, - status = FollowUpStatusType.Success, - ) - ), - ) - } - else -> { - when (err) { - is FeatureDevException -> { - messenger.sendError( - tabId = tabId, - errMessage = err.message, - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe - ) - } - else -> { - val msg = createUserFacingErrorMessage("$FEATURE_NAME request failed: ${err.message ?: err.cause?.message}") - val isDenyListedError = denyListedErrors.any { msg?.contains(it) ?: false } - val defaultMessage: String = when (session?.sessionState?.phase) { - SessionStatePhase.CODEGEN -> { - if (isDenyListedError || retriesRemaining(session) > 0) { - message("amazonqFeatureDev.code_generation.error_message") - } else { - message("amazonqFeatureDev.code_generation.no_retries.error_message") - } - } - else -> message("amazonqFeatureDev.error_text") - } - messenger.sendError( - tabId = tabId, - errMessage = defaultMessage, - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe - ) - } - } - } - } - } - - private suspend fun disablePreviousFileList(tabId: String) { - val session = getSessionInfo(tabId) - when (val sessionState = session.sessionState) { - is PrepareCodeGenerationState -> { - session.disableFileList(sessionState.filePaths, sessionState.deletedFiles, messenger) - } - } - } - - private fun storeCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) { - val tabId = message.tabId - val session = getSessionInfo(tabId) - session.storeCodeResultMessageId(message) - } - - private suspend fun handleChat( - tabId: String, - message: String, - ) { - var session: Session? = null - - this.disablePreviousFileList(tabId) - try { - logger.debug { "$FEATURE_NAME: Processing message: $message" } - session = getSessionInfo(tabId) - session.latestMessage = message - - val credentialState = authController.getAuthNeededStates(context.project).amazonQ - if (credentialState != null) { - messenger.sendAuthNeededException( - tabId = tabId, - triggerId = UUID.randomUUID().toString(), - credentialState = credentialState, - ) - session.isAuthenticating = true - return - } - - val codeWhispererSettings = CodeWhispererSettings.getInstance().getAutoBuildSetting() - val hasDevFile = session.context.hasDevFile() - val isPromptedForAutoBuildFeature = codeWhispererSettings.containsKey(session.context.workspaceRoot.path) - - if (hasDevFile && !isPromptedForAutoBuildFeature) { - promptAllowQCommandsConsent(messenger, tabId) - return - } - - session.preloader(messenger) - - when (session.sessionState.phase) { - SessionStatePhase.CODEGEN -> onCodeGeneration(session, message, tabId) - else -> null - } - } catch (err: Exception) { - processErrorChatMessage(err, message, session, tabId) - - // Lock the chat input until they explicitly click one of the follow-ups - messenger.sendChatInputEnabledMessage(tabId, enabled = false) - } - } - - private suspend fun promptAllowQCommandsConsent(messenger: MessagePublisher, tabID: String) { - messenger.sendAnswer( - tabId = tabID, - message = message("amazonqFeatureDev.chat_message.devFileInRepository"), - messageType = FeatureDevMessageType.Answer - ) - messenger.sendAnswer( - tabId = tabID, - messageType = FeatureDevMessageType.SystemPrompt, - followUp = listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.accept_for_project"), - type = FollowUpTypes.ACCEPT_AUTO_BUILD, - status = FollowUpStatusType.Success - ), - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.decline_for_project"), - type = FollowUpTypes.DENY_AUTO_BUILD, - status = FollowUpStatusType.Error - ) - ) - ) - } - - private suspend fun retryRequests(tabId: String) { - var session: Session? = null - try { - messenger.sendAsyncEventProgress( - tabId = tabId, - inProgress = true, - ) - session = getSessionInfo(tabId) - - // Decrease retries before making this request, just in case this one fails as well - session.decreaseRetries() - - // Sending an empty message will re-run the last state with the previous values - handleChat( - tabId = tabId, - message = session.latestMessage - ) - } catch (err: Exception) { - logger.error(err) { "Failed to retry request: ${err.message}" } - val message = createUserFacingErrorMessage("Failed to retry request: ${err.message}") - messenger.sendError( - tabId = tabId, - errMessage = message ?: message("amazonqFeatureDev.exception.retry_request_failed"), - retries = retriesRemaining(session), - conversationId = session?.conversationIdUnsafe, - ) - } finally { - // Finish processing the event - messenger.sendAsyncEventProgress( - tabId = tabId, - inProgress = false, - ) - } - } - - private suspend fun modifyDefaultSourceFolder(tabId: String) { - val session = getSessionInfo(tabId) - val workspaceRoot = session.context.workspaceRoot - - val modifyFolderFollowUp = FollowUp( - pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"), - type = FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, - status = FollowUpStatusType.Info, - ) - - var result: Result = Result.Failed - var reason: ModifySourceFolderErrorReason? = null - - withContext(EDT) { - val selectedFolder = selectFolder(context.project, workspaceRoot) - // No folder was selected - if (selectedFolder == null) { - logger.info { "Cancelled dialog and not selected any folder" } - - messenger.sendSystemPrompt( - tabId = tabId, - followUp = listOf(modifyFolderFollowUp), - ) - - reason = ModifySourceFolderErrorReason.ClosedBeforeSelection - return@withContext - } - - if (!selectedFolder.path.startsWith(workspaceRoot.path)) { - logger.info { "Selected folder not in workspace: ${selectedFolder.path}" } - - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.follow_up.incorrect_source_folder"), - ) - - messenger.sendSystemPrompt( - tabId = tabId, - followUp = listOf(modifyFolderFollowUp), - ) - - reason = ModifySourceFolderErrorReason.NotInWorkspaceFolder - return@withContext - } - - logger.info { "Selected correct folder inside workspace: ${selectedFolder.path}" } - - session.context.selectionRoot = selectedFolder - result = Result.Succeeded - - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.follow_up.modified_source_folder", selectedFolder.path), - canBeVoted = true, - ) - - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.SystemPrompt, - followUp = listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.retry"), - type = FollowUpTypes.RETRY, - status = FollowUpStatusType.Warning - ) - ), - ) - - messenger.sendChatInputEnabledMessage(tabId, enabled = false) - messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.write_new_prompt")) - } - - AmazonqTelemetry.modifySourceFolder( - amazonqConversationId = session.conversationId, - credentialStartUrl = getStartUrl(project = context.project), - result = result, - reason = reason?.toString() - ) - } - - private fun sendFeedback() { - runInEdt { - FeatureDevFeedbackDialog(context.project).show() - } - } - - fun getProject() = context.project - - private fun getSessionInfo(tabId: String) = chatSessionStorage.getSession(tabId, context.project) - - fun retriesRemaining(session: Session?): Int = session?.retries ?: DEFAULT_RETRY_LIMIT - - companion object { - private val logger = getLogger() - - private val objectMapper = jacksonObjectMapper() - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerExtensions.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerExtensions.kt deleted file mode 100644 index 0c1d926b992..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerExtensions.kt +++ /dev/null @@ -1,281 +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.amazonqFeatureDev.controller - -import com.intellij.notification.NotificationAction -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CLIENT_ERROR_MESSAGES -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CODE_GENERATION_RETRY_LIMIT -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ClientException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.LlmException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MetricDataOperationName -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MetricDataResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ServiceException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FeatureDevMessageType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpStatusType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpTypes -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswer -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendChatInputEnabledMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendCodeResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendSystemPrompt -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions -import software.aws.toolkits.jetbrains.utils.notifyInfo -import software.aws.toolkits.resources.message - -suspend fun FeatureDevController.onCodeGeneration( - session: Session, - message: String, - tabId: String, -) { - messenger.sendAsyncEventProgress( - tabId = tabId, - inProgress = true, - message = if (session.retries == CODE_GENERATION_RETRY_LIMIT) { - message( - "amazonqFeatureDev.chat_message.start_code_generation", - ) - } else { - message("amazonqFeatureDev.chat_message.start_code_generation_retry") - }, - ) - - try { - this.messenger.sendAnswer( - tabId = tabId, - message = message("amazonqFeatureDev.chat_message.requesting_changes"), - messageType = FeatureDevMessageType.AnswerStream, - ) - var state = session.sessionState - - var remainingIterations: Int? = state.codeGenerationRemainingIterationCount - var totalIterations: Int? = state.codeGenerationTotalIterationCount - - if (state.token?.token?.isCancellationRequested() == true) { - disposeToken(messenger, tabId, state.codeGenerationRemainingIterationCount, state.codeGenerationTotalIterationCount) - return - } - - messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.generating_code")) - - session.sendMetricDataTelemetry( - MetricDataOperationName.StartCodeGeneration, - MetricDataResult.Success - ) - - session.send(message) // Trigger code generation - - state = session.sessionState - - var filePaths: List = emptyList() - var deletedFiles: List = emptyList() - var references: List = emptyList() - var uploadId = "" - - when (state) { - is PrepareCodeGenerationState -> { - filePaths = state.filePaths - deletedFiles = state.deletedFiles - references = state.references - uploadId = state.uploadId - remainingIterations = state.codeGenerationRemainingIterationCount - totalIterations = state.codeGenerationTotalIterationCount - } - } - - if (state.token?.token?.isCancellationRequested() == true) { - disposeToken(messenger, tabId, state.codeGenerationRemainingIterationCount, state.codeGenerationTotalIterationCount) - return - } - - // Atm this is the only possible path as codegen is mocked to return empty. - if (filePaths.isEmpty() && deletedFiles.isEmpty()) { - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.code_generation.no_file_changes"), - ) - messenger.sendSystemPrompt( - tabId = tabId, - followUp = - if (retriesRemaining(session) > 0) { - listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.retry"), - type = FollowUpTypes.RETRY, - status = FollowUpStatusType.Warning, - ), - ) - } else { - emptyList() - }, - ) - messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until retry is clicked. - return - } - - messenger.sendCodeResult(tabId = tabId, uploadId = uploadId, filePaths = filePaths, deletedFiles = deletedFiles, references = references) - - if (remainingIterations != null && totalIterations != null) { - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = - if (remainingIterations > 2) { - message("amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code_or_feedback") - } else if (remainingIterations > 0) { - message( - "amazonqFeatureDev.code_generation.iteration_counts", - remainingIterations, - totalIterations, - ) - } else { - message( - "amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code", - remainingIterations, - totalIterations, - ) - }, - ) - } - - messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase, InsertAction.ALL)) - messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation")) - } catch (err: Exception) { - val metricDataResult: MetricDataResult - when (err) { - is ClientException, - -> { - metricDataResult = MetricDataResult.Error - } - - is LlmException -> { - metricDataResult = MetricDataResult.LlmFailure - } - - is ServiceException -> { - metricDataResult = MetricDataResult.Fault - } - - else -> { - val errorMessage = err.message.orEmpty() - metricDataResult = if (CLIENT_ERROR_MESSAGES.any { errorMessage.contains(it) }) { - MetricDataResult.Error - } else { - MetricDataResult.Fault - } - } - } - session.sendMetricDataTelemetry( - MetricDataOperationName.EndCodeGeneration, - metricDataResult - ) - throw err - } finally { - if (session.sessionState.token - ?.token - ?.isCancellationRequested() == true - ) { - session.sessionState.token = CancellationTokenSource() - } else { - messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) // Finish processing the event - messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until a follow-up is clicked. - } - if (toolWindow != null && !toolWindow.isVisible) { - notifyInfo( - title = message("amazonqFeatureDev.code_generation.notification_title"), - content = message("amazonqFeatureDev.code_generation.notification_message"), - project = getProject(), - notificationActions = listOf(openChatNotificationAction()), - ) - } - } - - session.sendMetricDataTelemetry( - MetricDataOperationName.EndCodeGeneration, - MetricDataResult.Success - ) -} - -private suspend fun disposeToken( - messenger: MessagePublisher, - tabId: String, - remainingIterations: Int?, - totalIterations: Int?, -) { - if (remainingIterations !== null && remainingIterations <= 0) { - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = - message( - "amazonqFeatureDev.code_generation.stopped_code_generation_no_iterations", - ), - ) - // I stopped generating your code. You don't have more iterations left, however, you can start a new session - messenger.sendSystemPrompt( - tabId = tabId, - followUp = - listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.new_task"), - type = FollowUpTypes.NEW_TASK, - status = FollowUpStatusType.Info, - ), - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.close_session"), - type = FollowUpTypes.CLOSE_SESSION, - status = FollowUpStatusType.Info, - ), - ), - ) - messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) - messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation")) - - return - } - - if (remainingIterations !== null && totalIterations !== null && remainingIterations <= 2) { - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = - message( - "amazonqFeatureDev.code_generation.stopped_code_generation", - remainingIterations, - totalIterations, - ), - ) - } else { - messenger.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = - message("amazonqFeatureDev.code_generation.stopped_code_generation_no_iteration_count_display"), - ) - } - - messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true) - - messenger.sendUpdatePlaceholder( - tabId = tabId, - newPlaceholder = message("amazonqFeatureDev.placeholder.new_plan"), - ) -} - -private fun FeatureDevController.openChatNotificationAction() = - NotificationAction.createSimple( - message("amazonqFeatureDev.code_generation.notification_open_link"), - ) { - toolWindow?.show() - } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessage.kt deleted file mode 100644 index 9d82818c4a9..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessage.kt +++ /dev/null @@ -1,261 +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.amazonqFeatureDev.messages - -import com.fasterxml.jackson.annotation.JsonProperty -import com.fasterxml.jackson.annotation.JsonValue -import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthFollowUpType -import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo -import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference -import java.time.Instant -import java.util.UUID - -sealed interface FeatureDevBaseMessage : AmazonQMessage - -// === UI -> App Messages === -sealed interface IncomingFeatureDevMessage : FeatureDevBaseMessage { - - data class ChatPrompt( - val chatMessage: String, - val command: String, - @JsonProperty("tabID") val tabId: String, - ) : IncomingFeatureDevMessage - - data class StoreMessageIdMessage( - @JsonProperty("tabID") val tabId: String, - val command: String, - val messageId: String?, - ) : IncomingFeatureDevMessage - - data class NewTabCreated( - val command: String, - @JsonProperty("tabID") val tabId: String, - ) : IncomingFeatureDevMessage - - data class AuthFollowUpWasClicked( - @JsonProperty("tabID") val tabId: String, - val authType: AuthFollowUpType, - ) : IncomingFeatureDevMessage - - data class TabRemoved( - val command: String, - @JsonProperty("tabID") val tabId: String, - ) : IncomingFeatureDevMessage - - data class FollowupClicked( - val followUp: FollowUp, - @JsonProperty("tabID") val tabId: String, - val messageId: String?, - val command: String, - ) : IncomingFeatureDevMessage - - data class ChatItemVotedMessage( - @JsonProperty("tabID") val tabId: String, - val messageId: String, - val vote: String, - ) : IncomingFeatureDevMessage - - data class ChatItemFeedbackMessage( - @JsonProperty("tabID") val tabId: String, - val selectedOption: String, - val comment: String?, - val messageId: String, - ) : IncomingFeatureDevMessage - - data class ClickedLink( - @JsonProperty("tabID") val tabId: String, - val command: String, - val messageId: String?, - val link: String, - ) : IncomingFeatureDevMessage - - data class StopResponse( - @JsonProperty("tabID") val tabId: String, - val command: String, - val messageId: String?, - ) : IncomingFeatureDevMessage - - data class InsertCodeAtCursorPosition( - @JsonProperty("tabID") val tabId: String, - val code: String, - val insertionTargetType: String?, - val codeReference: List?, - ) : IncomingFeatureDevMessage - - data class OpenDiff( - @JsonProperty("tabID") val tabId: String, - val filePath: String, - val deleted: Boolean, - ) : IncomingFeatureDevMessage - - data class FileClicked( - @JsonProperty("tabID") val tabId: String, - val filePath: String, - val messageId: String, - val actionName: String, - ) : IncomingFeatureDevMessage -} - -// === UI -> App Messages === - -sealed class UiMessage( - open val tabId: String?, - open val type: String, -) : FeatureDevBaseMessage { - val time = Instant.now().epochSecond - val sender = "featureDevChat" -} - -enum class FeatureDevMessageType( - @field:JsonValue val json: String, -) { - Answer("answer"), - AnswerPart("answer-part"), - AnswerStream("answer-stream"), - SystemPrompt("system-prompt"), -} - -data class FeatureDevMessage( - @JsonProperty("tabID") override val tabId: String, - @JsonProperty("triggerID") val triggerId: String, - val messageType: FeatureDevMessageType, - val messageId: String, - val message: String? = null, - val followUps: List? = null, - val canBeVoted: Boolean, - val snapToTop: Boolean, - -) : UiMessage( - tabId = tabId, - type = "chatMessage", -) - -data class AsyncEventProgressMessage( - @JsonProperty("tabID") override val tabId: String, - val message: String? = null, - val inProgress: Boolean, -) : UiMessage( - tabId = tabId, - type = "asyncEventProgressMessage" -) - -data class UpdatePlaceholderMessage( - @JsonProperty("tabID") override val tabId: String, - val newPlaceholder: String, -) : UiMessage( - tabId = tabId, - type = "updatePlaceholderMessage" -) - -data class FileComponent( - @JsonProperty("tabID") override val tabId: String, - val filePaths: List, - val deletedFiles: List, - val messageId: String, - val disableFileActions: Boolean = false, -) : UiMessage( - tabId = tabId, - type = "updateFileComponent" -) - -data class ChatInputEnabledMessage( - @JsonProperty("tabID") override val tabId: String, - val enabled: Boolean, -) : UiMessage( - tabId = tabId, - type = "chatInputEnabledMessage" -) -data class ErrorMessage( - @JsonProperty("tabID") override val tabId: String, - val title: String, - val message: String, -) : UiMessage( - tabId = tabId, - type = "errorMessage", -) - -data class AuthenticationUpdateMessage( - val authenticatingTabIDs: List, - val featureDevEnabled: Boolean, - val codeTransformEnabled: Boolean, - val codeScanEnabled: Boolean, - val codeTestEnabled: Boolean, - val docEnabled: Boolean, - val message: String? = null, - val messageId: String = UUID.randomUUID().toString(), - -) : UiMessage( - null, - type = "authenticationUpdateMessage", -) - -data class AuthNeededException( - @JsonProperty("tabID") override val tabId: String, - @JsonProperty("triggerID") val triggerId: String, - val authType: AuthFollowUpType, - val message: String, -) : UiMessage( - tabId = tabId, - type = "authNeededException", -) - -data class CodeResultMessage( - @JsonProperty("tabID") override val tabId: String, - val conversationId: String, - val filePaths: List, - val deletedFiles: List, - val references: List, - val messageId: String?, -) : UiMessage( - tabId = tabId, - type = "codeResultMessage" -) - -data class FollowUp( - val type: FollowUpTypes, - val pillText: String, - val disabled: Boolean? = false, - val description: String? = null, - val status: FollowUpStatusType? = null, - val icon: FollowUpIcons? = null, -) - -enum class FollowUpIcons( - @field:JsonValue val json: String, -) { - Ok("ok"), - Refresh("refresh"), -} - -enum class FollowUpStatusType( - @field:JsonValue val json: String, -) { - Info("info"), - Success("success"), - Warning("warning"), - Error("error"), -} - -enum class FollowUpTypes( - @field:JsonValue val json: String, -) { - RETRY("Retry"), - MODIFY_DEFAULT_SOURCE_FOLDER("ModifyDefaultSourceFolder"), - DEV_EXAMPLES("DevExamples"), - SEND_FEEDBACK("SendFeedback"), - INSERT_CODE("InsertCode"), - PROVIDE_FEEDBACK_AND_REGENERATE_CODE("ProvideFeedbackAndRegenerateCode"), - NEW_TASK("NewTask"), - CLOSE_SESSION("CloseSession"), - ACCEPT_AUTO_BUILD("AcceptAutoBuild"), - DENY_AUTO_BUILD("DenyAutoBuild"), - GENERATE_DEV_FILE("GenerateDevFile"), -} - -// Util classes -data class ReducedCodeReference( - val information: String, -) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessagePublisherExtensions.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessagePublisherExtensions.kt deleted file mode 100644 index 5e757bdef60..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessagePublisherExtensions.kt +++ /dev/null @@ -1,221 +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.amazonqFeatureDev.messages - -import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo -import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference -import software.aws.toolkits.jetbrains.services.cwc.messages.RecommendationContentSpan -import software.aws.toolkits.resources.message -import java.util.UUID - -suspend fun MessagePublisher.sendAnswer( - tabId: String, - message: String? = null, - messageId: String? = null, - messageType: FeatureDevMessageType, - followUp: List? = null, - canBeVoted: Boolean? = false, - snapToTop: Boolean? = false, -) { - val chatMessage = - FeatureDevMessage( - tabId = tabId, - triggerId = UUID.randomUUID().toString(), - messageId = messageId ?: UUID.randomUUID().toString(), - messageType = messageType, - message = message, - followUps = followUp, - canBeVoted = canBeVoted ?: false, - snapToTop = snapToTop ?: false, - ) - this.publish(chatMessage) -} - -suspend fun MessagePublisher.sendAnswerPart( - tabId: String, - message: String? = null, - canBeVoted: Boolean? = null, -) { - this.sendAnswer( - tabId = tabId, - message = message, - messageType = FeatureDevMessageType.AnswerPart, - canBeVoted = canBeVoted, - ) -} - -/** Send or replace system prompt. Only one system prompt can be rendered for a given tab. */ -suspend fun MessagePublisher.sendSystemPrompt( - tabId: String, - followUp: List, -) { - this.sendAnswer( - tabId = tabId, - messageId = "$tabId-system-prompt", - messageType = FeatureDevMessageType.SystemPrompt, - followUp = followUp, - ) -} - -suspend fun MessagePublisher.updateFileComponent( - tabId: String, - filePaths: List, - deletedFiles: List, - messageId: String, - disableFileActions: Boolean = false, -) { - val fileComponentMessage = FileComponent( - tabId = tabId, - filePaths = filePaths, - deletedFiles = deletedFiles, - messageId = messageId, - disableFileActions = disableFileActions, - ) - this.publish(fileComponentMessage) -} - -suspend fun MessagePublisher.sendAsyncEventProgress( - tabId: String, - inProgress: Boolean, - message: String? = null, -) { - val asyncEventProgressMessage = AsyncEventProgressMessage( - tabId = tabId, - message = message, - inProgress = inProgress, - ) - this.publish(asyncEventProgressMessage) -} - -suspend fun MessagePublisher.sendUpdatePlaceholder(tabId: String, newPlaceholder: String) { - val updatePlaceholderMessage = UpdatePlaceholderMessage( - tabId = tabId, - newPlaceholder = newPlaceholder, - ) - this.publish(updatePlaceholderMessage) -} - -suspend fun MessagePublisher.sendAuthNeededException(tabId: String, triggerId: String, credentialState: AuthNeededState) { - val message = AuthNeededException( - tabId = tabId, - triggerId = triggerId, - authType = credentialState.authType, - message = credentialState.message, - ) - this.publish(message) -} - -suspend fun MessagePublisher.sendAuthenticationInProgressMessage(tabId: String) { - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.follow_instructions_for_authentication"), - ) -} - -suspend fun MessagePublisher.sendChatInputEnabledMessage(tabId: String, enabled: Boolean) { - val chatInputEnabledMessage = ChatInputEnabledMessage( - tabId, - enabled, - ) - this.publish(chatInputEnabledMessage) -} - -suspend fun MessagePublisher.sendError(tabId: String, errMessage: String?, retries: Int, conversationId: String? = null, showDefaultMessage: Boolean? = false) { - val conversationIdText = if (conversationId == null) "" else "\n\nConversation ID: **$conversationId**" - - if (retries == 0) { - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = if (showDefaultMessage == true) errMessage else message("amazonqFeatureDev.no_retries.error_text") + conversationIdText, - ) - - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.SystemPrompt, - followUp = listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.send_feedback"), - type = FollowUpTypes.SEND_FEEDBACK, - status = FollowUpStatusType.Info, - ), - ), - ) - return - } - - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = errMessage + conversationIdText, - ) - - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.SystemPrompt, - followUp = - listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.retry"), - type = FollowUpTypes.RETRY, - status = FollowUpStatusType.Warning, - ), - ), - ) -} - -suspend fun MessagePublisher.sendMonthlyLimitError(tabId: String) { - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.exception.monthly_limit_error"), - ) - this.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_monthly_limit")) -} - -suspend fun MessagePublisher.initialExamples(tabId: String) { - this.sendAnswer( - tabId = tabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.example_text"), - ) -} - -suspend fun MessagePublisher.sendCodeResult( - tabId: String, - uploadId: String, - filePaths: List, - deletedFiles: List, - references: List, -) { - val messageId = UUID.randomUUID() - val refs = references.map { ref -> - CodeReference( - licenseName = ref.licenseName, - repository = ref.repository, - url = ref.url, - recommendationContentSpan = RecommendationContentSpan( - ref.recommendationContentSpan?.start ?: 0, - ref.recommendationContentSpan?.end ?: 0, - ), - information = "Reference code under **${ref.licenseName}** license from repository [${ref.repository}](${ref.url})", - ) - } - - this.publish( - CodeResultMessage( - tabId = tabId, - conversationId = uploadId, - filePaths = filePaths, - deletedFiles = deletedFiles, - references = refs, - messageId = messageId.toString(), - ), - ) -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt deleted file mode 100644 index 9647ca9e359..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt +++ /dev/null @@ -1,308 +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.amazonqFeatureDev.session - -import kotlinx.coroutines.delay -import software.amazon.awssdk.services.codewhispererruntime.model.CodeGenerationWorkflowStatus -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeGenerationException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.EmptyPatchException -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.FileCreationFailedException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.GuardrailsException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.NoChangeRequiredException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.PromptRefusalException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ThrottlingException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getChangeIdentifier -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getDiffMetrics -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.readFileToString -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl -import software.aws.toolkits.resources.message -import software.aws.toolkits.telemetry.AmazonqTelemetry -import software.aws.toolkits.telemetry.MetricResult -import java.util.UUID - -private val logger = getLogger() -private const val RUN_COMMAND_LOG_PATH = ".amazonq/dev/run_command.log" - -class CodeGenerationState( - override val tabID: String, - override var approach: String, - val config: SessionStateConfig, - val uploadId: String, - override var currentIteration: Int? = 0, - val repositorySize: Double, - val messenger: MessagePublisher, - override var codeGenerationRemainingIterationCount: Int? = null, - override var codeGenerationTotalIterationCount: Int? = null, - var currentCodeGenerationId: String? = "EMPTY_CURRENT_CODE_GENERATION_ID", - override var token: CancellationTokenSource?, - override var diffMetricsProcessed: DiffMetricsProcessed, -) : SessionState { - override val phase = SessionStatePhase.CODEGEN - - override suspend fun interact(action: SessionStateAction): SessionStateInteraction { - val startTime = System.currentTimeMillis() - var result: MetricResult = MetricResult.Succeeded - var failureReason: String? = null - var failureReasonDesc: String? = null - var codeGenerationWorkflowStatus: CodeGenerationWorkflowStatus = CodeGenerationWorkflowStatus.COMPLETE - var numberOfReferencesGenerated: Int? = null - var numberOfFilesGenerated: Int? = null - try { - val codeGenerationId = UUID.randomUUID() - - val response = - config.featureDevService.startTaskAssistCodeGeneration( - conversationId = config.conversationId, - uploadId = uploadId, - message = action.msg, - codeGenerationId = codeGenerationId.toString(), - currentCodeGenerationId = currentCodeGenerationId.toString(), - ) - - if (action.token?.token?.isCancellationRequested() != true) { - this.currentCodeGenerationId = codeGenerationId.toString() - } - - messenger.sendAnswerPart( - tabId = tabID, - message = message("amazonqFeatureDev.code_generation.generating_code"), - ) - messenger.sendUpdatePlaceholder( - tabId = tabID, - newPlaceholder = message("amazonqFeatureDev.code_generation.generating_code"), - ) - val codeGenerationResult = generateCode(codeGenerationId = response.codeGenerationId(), messenger = messenger, token = action.token) - numberOfReferencesGenerated = codeGenerationResult.references.size - numberOfFilesGenerated = codeGenerationResult.newFiles.size - codeGenerationRemainingIterationCount = codeGenerationResult.codeGenerationRemainingIterationCount - codeGenerationTotalIterationCount = codeGenerationResult.codeGenerationTotalIterationCount - currentIteration = - if (codeGenerationRemainingIterationCount != null && codeGenerationTotalIterationCount != null) { - codeGenerationTotalIterationCount?.let { total -> codeGenerationRemainingIterationCount?.let { remaining -> total - remaining } } - } else { - currentIteration?.plus(1) - } - - runCatching { - var insertedLines = 0 - var insertedCharacters = 0 - codeGenerationResult.newFiles.forEach { file -> - // FIXME: Ideally, the before content should be read from the uploaded context instead of from disk, to avoid drift - val before = config.repoContext.addressableRoot - .toNioPath() - .resolve(file.zipFilePath) - .toFile() - .let { f -> - if (f.exists() && f.canRead()) { - readFileToString(f) - } else { - "" - } - } - - val changeIdentifier = getChangeIdentifier(file.zipFilePath, before, file.fileContent) - - if (!diffMetricsProcessed.generated.contains(changeIdentifier)) { - val diffMetrics = getDiffMetrics(before, file.fileContent) - insertedLines += diffMetrics.insertedLines - insertedCharacters += diffMetrics.insertedCharacters - diffMetricsProcessed.generated.add(changeIdentifier) - } - } - if (insertedLines > 0) { - config.featureDevService.sendFeatureDevCodeGenerationEvent( - conversationId = config.conversationId, - linesOfCodeGenerated = insertedLines, - charactersOfCodeGenerated = insertedCharacters, - ) - } - }.onFailure { /* Noop on diff telemetry failure */ } - - val nextState = - PrepareCodeGenerationState( - tabID = tabID, - approach = approach, - config = config, - filePaths = codeGenerationResult.newFiles, - deletedFiles = codeGenerationResult.deletedFiles, - references = codeGenerationResult.references, - currentIteration = currentIteration, - uploadId = uploadId, - messenger = messenger, - codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount, - codeGenerationTotalIterationCount = codeGenerationTotalIterationCount, - token = this.token, - diffMetricsProcessed = diffMetricsProcessed, - ) - - // It is not needed to interact right away with the PrepareCodeGeneration. - // returns therefore a SessionStateInteraction object to be handled by the controller. - return SessionStateInteraction( - nextState = nextState, - interaction = Interaction(content = "", interactionSucceeded = true), - ) - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: Code generation failed: ${e.message}" } - result = MetricResult.Failed - failureReason = e.javaClass.simpleName - if (e is FeatureDevException) { - failureReason = e.reason() - failureReasonDesc = e.reasonDesc() - } - codeGenerationWorkflowStatus = CodeGenerationWorkflowStatus.FAILED - - throw e - } finally { - currentIteration?.let { - AmazonqTelemetry.codeGenerationInvoke( - amazonqConversationId = config.conversationId, - amazonqCodeGenerationResult = codeGenerationWorkflowStatus.toString(), - amazonqGenerateCodeIteration = it.toDouble(), - amazonqNumberOfReferences = numberOfReferencesGenerated?.toDouble(), - amazonqGenerateCodeResponseLatency = (System.currentTimeMillis() - startTime).toDouble(), - amazonqNumberOfFilesGenerated = numberOfFilesGenerated?.toDouble(), - amazonqRepositorySize = repositorySize, - result = result, - reason = failureReason, - reasonDesc = failureReasonDesc, - duration = (System.currentTimeMillis() - startTime).toDouble(), - credentialStartUrl = getStartUrl(config.featureDevService.project), - ) - } - } - } -} - -private suspend fun CodeGenerationState.generateCode( - codeGenerationId: String, - messenger: MessagePublisher, - token: CancellationTokenSource?, -): CodeGenerationResult { - val pollCount = 360 - val requestDelay = 5000L - var codeGenerationRemainingIterationCount: Int? = null - var codeGenerationTotalIterationCount: Int? = null - - repeat(pollCount) { - if (token?.token?.isCancellationRequested() == true) { - return CodeGenerationResult(emptyList(), emptyList(), emptyList(), codeGenerationRemainingIterationCount, codeGenerationTotalIterationCount) - } - val codeGenerationResultState = - config.featureDevService.getTaskAssistCodeGeneration( - conversationId = config.conversationId, - codeGenerationId = codeGenerationId, - ) - - codeGenerationRemainingIterationCount = codeGenerationResultState.codeGenerationRemainingIterationCount() - codeGenerationTotalIterationCount = codeGenerationResultState.codeGenerationTotalIterationCount() - - when (codeGenerationResultState.codeGenerationStatus().status()) { - CodeGenerationWorkflowStatus.COMPLETE -> { - val codeGenerationStreamResult = - config.featureDevService.exportTaskAssistArchiveResult( - conversationId = config.conversationId, - ) - - val fileContents = codeGenerationStreamResult.new_file_contents.filterKeys { file -> - if (file.endsWith(RUN_COMMAND_LOG_PATH)) { - val contents: String = codeGenerationStreamResult.new_file_contents[file].orEmpty() - val truncatedContents = if (contents.length > 10000000) { - contents.substring(0, 10000000) - } else { - contents - } - logger.info(truncatedContents) { "Run command log: $truncatedContents" } - false - } else { - true - } - } - - val newFileInfo = registerNewFiles(newFileContents = fileContents) - val deletedFileInfo = registerDeletedFiles(deletedFiles = codeGenerationStreamResult.deleted_files) - return CodeGenerationResult( - newFiles = newFileInfo, - deletedFiles = deletedFileInfo, - references = codeGenerationStreamResult.references, - codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount, - codeGenerationTotalIterationCount = codeGenerationTotalIterationCount, - ) - } - CodeGenerationWorkflowStatus.IN_PROGRESS -> { - if (codeGenerationResultState.codeGenerationStatusDetail() != null) { - messenger.sendAnswerPart( - tabId = tabID, - message = - message("amazonqFeatureDev.code_generation.generating_code") + - "\n\n" + codeGenerationResultState.codeGenerationStatusDetail(), - ) - } - delay(requestDelay) - } - CodeGenerationWorkflowStatus.FAILED -> { - when (true) { - codeGenerationResultState.codeGenerationStatusDetail()?.contains( - "Guardrails", - ), - -> throw GuardrailsException(operation = FeatureDevOperation.GenerateCode.toString(), desc = "Failed guardrails") - codeGenerationResultState.codeGenerationStatusDetail()?.contains( - "PromptRefusal", - ), - -> throw PromptRefusalException(operation = FeatureDevOperation.GenerateCode.toString(), desc = "Prompt refusal") - codeGenerationResultState.codeGenerationStatusDetail()?.contains( - "EmptyPatch", - ), - -> throw FileCreationFailedException(operation = FeatureDevOperation.GenerateCode.toString(), desc = "File creation failed") - codeGenerationResultState.codeGenerationStatusDetail()?.contains( - "FileCreationFailed", - ), - -> { - if (codeGenerationResultState.codeGenerationStatusDetail().contains("NO_CHANGE_REQUIRED")) { - throw NoChangeRequiredException(operation = FeatureDevOperation.GenerateCode.toString(), desc = "No change required") - } - throw EmptyPatchException(operation = FeatureDevOperation.GenerateCode.toString(), desc = "Empty patch") - } - codeGenerationResultState.codeGenerationStatusDetail()?.contains( - "Throttling", - ), - -> throw ThrottlingException(operation = FeatureDevOperation.GenerateCode.toString(), desc = "Request throttled") - else -> throw CodeGenerationException(operation = FeatureDevOperation.GenerateCode.toString(), desc = null) - } - } - else -> error("Unknown status: ${codeGenerationResultState.codeGenerationStatus().status()}") - } - } - - return CodeGenerationResult(emptyList(), emptyList(), emptyList(), codeGenerationRemainingIterationCount, codeGenerationTotalIterationCount) -} - -fun registerNewFiles(newFileContents: Map): List = - newFileContents.map { - NewFileZipInfo( - // Note: When managing file state, we normalize file paths returned from the agent in order to ensure they are handled as relative paths. - zipFilePath = it.key.removePrefix("/"), - fileContent = it.value, - rejected = false, - changeApplied = false - ) - } - -fun registerDeletedFiles(deletedFiles: List): List = - deletedFiles.map { - DeletedFileInfo( - // Note: When managing file state, we normalize file paths returned from the agent in order to ensure they are handled as relative paths. - zipFilePath = it.removePrefix("/"), - rejected = false, - changeApplied = false - ) - } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/ConversationNotStartedState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/ConversationNotStartedState.kt deleted file mode 100644 index 7bfa85b39a6..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/ConversationNotStartedState.kt +++ /dev/null @@ -1,22 +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.amazonqFeatureDev.session - -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource - -class ConversationNotStartedState( - override var approach: String, - override val tabID: String, - override var token: CancellationTokenSource?, - override var codeGenerationRemainingIterationCount: Int? = null, - override var codeGenerationTotalIterationCount: Int? = null, - override var currentIteration: Int?, - override var diffMetricsProcessed: DiffMetricsProcessed, -) : 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/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt deleted file mode 100644 index 7ffdfe67ebc..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt +++ /dev/null @@ -1,104 +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.amazonqFeatureDev.session - -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.deleteUploadArtifact -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3 -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl -import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings -import software.aws.toolkits.resources.message -import software.aws.toolkits.telemetry.AmazonqTelemetry -import software.aws.toolkits.telemetry.AmazonqUploadIntent -import software.aws.toolkits.telemetry.Result -import java.util.UUID - -private val logger = getLogger() - -class PrepareCodeGenerationState( - override var tabID: String, - override var token: CancellationTokenSource?, - override var approach: String, - private var config: SessionStateConfig, - val filePaths: List, - val deletedFiles: List, - val references: List, - var uploadId: String, - override var currentIteration: Int?, - private var messenger: MessagePublisher, - override var codeGenerationRemainingIterationCount: Int? = null, - override var codeGenerationTotalIterationCount: Int? = null, - override var diffMetricsProcessed: DiffMetricsProcessed, -) : SessionState { - override val phase = SessionStatePhase.CODEGEN - override suspend fun interact(action: SessionStateAction): SessionStateInteraction { - val startTime = System.currentTimeMillis() - var result: Result = Result.Succeeded - var failureReason: String? = null - var failureReasonDesc: String? = null - var zipFileLength: Long? = null - val nextState: SessionState - try { - messenger.sendAnswerPart(tabId = this.tabID, message = message("amazonqFeatureDev.chat_message.uploading_code")) - messenger.sendUpdatePlaceholder(tabId = this.tabID, newPlaceholder = message("amazonqFeatureDev.chat_message.uploading_code")) - - val isAutoBuildFeatureEnabled = CodeWhispererSettings.getInstance().isAutoBuildFeatureEnabled(this.config.repoContext.workspaceRoot.path) - val repoZipResult = config.repoContext.getProjectZip(isAutoBuildFeatureEnabled = isAutoBuildFeatureEnabled) - val zipFileChecksum = repoZipResult.checksum - zipFileLength = repoZipResult.contentLength - val fileToUpload = repoZipResult.payload - - val uploadId = UUID.randomUUID() - val uploadUrlResponse = config.featureDevService.createUploadUrl( - config.conversationId, - zipFileChecksum, - zipFileLength, - uploadId.toString() - ) - - uploadArtifactToS3(uploadUrlResponse.uploadUrl(), fileToUpload, zipFileChecksum, zipFileLength, uploadUrlResponse.kmsKeyArn()) - deleteUploadArtifact(fileToUpload) - - this.uploadId = uploadId.toString() - messenger.sendAnswerPart(tabId = this.tabID, message = message("amazonqFeatureDev.placeholder.context_gathering_complete")) - messenger.sendUpdatePlaceholder(tabId = this.tabID, newPlaceholder = message("amazonqFeatureDev.placeholder.context_gathering_complete")) - nextState = CodeGenerationState( - tabID = this.tabID, - approach = "", // No approach needed, - config = this.config, - uploadId = this.uploadId, - currentIteration = this.currentIteration, - repositorySize = zipFileLength.toDouble(), - messenger = messenger, - token = this.token, - diffMetricsProcessed = diffMetricsProcessed - ) - } catch (e: Exception) { - result = Result.Failed - failureReason = e.javaClass.simpleName - failureReasonDesc = e.message - logger.warn(e) { "$FEATURE_NAME: Code uploading failed: ${e.message}" } - throw e - } finally { - AmazonqTelemetry.createUpload( - amazonqConversationId = config.conversationId, - amazonqRepositorySize = zipFileLength?.toDouble(), - amazonqUploadIntent = AmazonqUploadIntent.TASKASSISTPLANNING, - result = result, - reason = failureReason, - reasonDesc = failureReasonDesc, - duration = (System.currentTimeMillis() - startTime).toDouble(), - credentialStartUrl = getStartUrl(config.featureDevService.project) - ) - } - // It is essential to interact with the next state outside try-catch block for the telemetry to capture events for the states separately - return nextState.interact(action) - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt deleted file mode 100644 index c0ab7fc72a2..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt +++ /dev/null @@ -1,288 +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.amazonqFeatureDev.session - -import com.intellij.openapi.diagnostic.logger -import com.intellij.openapi.project.Project -import com.intellij.openapi.vfs.VfsUtil -import software.aws.toolkits.jetbrains.common.util.resolveAndCreateOrUpdateFile -import software.aws.toolkits.jetbrains.common.util.resolveAndDeleteFile -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CODE_GENERATION_RETRY_LIMIT -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ConversationIdNotFoundException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MAX_PROJECT_SIZE_BYTES -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MetricDataOperationName -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MetricDataResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.IncomingFeatureDevMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateFileComponent -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getChangeIdentifier -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getDiffMetrics -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.readFileToString -import software.aws.toolkits.jetbrains.services.cwc.controller.ReferenceLogController -import java.util.HashSet - -class Session(val tabID: String, val project: Project) { - var context: FeatureDevSessionContext - val sessionStartTime = System.currentTimeMillis() - - private var _state: SessionState? - private var preloaderFinished: Boolean = false - private var _conversationId: String? = null - private var _latestMessage: String = "" - private var task: String = "" - private val proxyClient: FeatureDevClient - private val featureDevService: FeatureDevService - private var _codeResultMessageId: String? = null - - // retry session state vars - private var codegenRetries: Int - - // Used to keep track of whether the current session/tab is currently authenticating/needs authenticating - var isAuthenticating: Boolean - - init { - context = FeatureDevSessionContext(project, MAX_PROJECT_SIZE_BYTES) - proxyClient = FeatureDevClient.getInstance(project) - featureDevService = FeatureDevService(proxyClient, project) - _state = ConversationNotStartedState( - approach = "", - tabID = tabID, - token = null, - currentIteration = 0, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()) - ) - isAuthenticating = false - codegenRetries = CODE_GENERATION_RETRY_LIMIT - } - - fun conversationIDLog(conversationId: String) = "$FEATURE_NAME Conversation ID: $conversationId" - - /** - * Preload any events that have to run before a chat message can be sent - */ - suspend fun preloader(messenger: MessagePublisher) { - if (!preloaderFinished) { - setupConversation(messenger) - preloaderFinished = true - messenger.sendAsyncEventProgress(tabId = this.tabID, inProgress = true) - featureDevService.sendFeatureDevEvent(this.conversationId) - } - } - - /** - * Starts a conversation with the backend and uploads the repo for the LLMs to be able to use it. - */ - private fun setupConversation(messenger: MessagePublisher) { - _conversationId = featureDevService.createConversation() - logger().info(conversationIDLog(this.conversationId)) - - val sessionStateConfig = getSessionStateConfig().copy(conversationId = this.conversationId) - _state = PrepareCodeGenerationState( - tabID = sessionState.tabID, - approach = sessionState.approach, - config = sessionStateConfig, - filePaths = emptyList(), - deletedFiles = emptyList(), - references = emptyList(), - currentIteration = 1, // first code gen iteration - uploadId = "", // There is no code gen uploadId so far - messenger = messenger, - token = CancellationTokenSource(), - diffMetricsProcessed = sessionState.diffMetricsProcessed, - ) - } - - fun storeCodeResultMessageId(message: IncomingFeatureDevMessage.StoreMessageIdMessage) { - val messageId = message.messageId - this.updateCodeResultMessageId(messageId) - } - - private fun updateCodeResultMessageId(messageId: String?) { - this._codeResultMessageId = messageId - } - - suspend fun updateFilesPaths( - filePaths: List, - deletedFiles: List, - messenger: MessagePublisher, - disableFileActions: Boolean = false, - ) { - val codeResultMessageId = this._codeResultMessageId - if (codeResultMessageId != null) { - messenger.updateFileComponent(this.tabID, filePaths, deletedFiles, codeResultMessageId, disableFileActions) - } - } - - /** - * Triggered by the Insert code follow-up button to apply code changes. - */ - suspend fun insertChanges( - filePaths: List, - deletedFiles: List, - references: List, - ) { - val newFilePaths = filePaths.filter { !it.rejected && !it.changeApplied } - val newDeletedFiles = deletedFiles.filter { !it.rejected && !it.changeApplied } - - runCatching { - var insertedLines = 0 - var insertedCharacters = 0 - filePaths.forEach { file -> - // FIXME: Ideally, the before content should be read from the uploaded context instead of from disk, to avoid drift - val before = context.addressableRoot.toNioPath() - .resolve(file.zipFilePath) - .toFile() - .let { f -> - if (f.exists() && f.canRead()) { - readFileToString(f) - } else { - "" - } - } - - val changeIdentifier = getChangeIdentifier(file.zipFilePath, before, file.fileContent) - - if (_state?.diffMetricsProcessed?.accepted?.contains(changeIdentifier) != true) { - val diffMetrics = getDiffMetrics(before, file.fileContent) - insertedLines += diffMetrics.insertedLines - insertedCharacters += diffMetrics.insertedCharacters - _state?.diffMetricsProcessed?.accepted?.add(changeIdentifier) - } - } - - if (insertedLines > 0) { - featureDevService.sendFeatureDevCodeAcceptanceEvent( - conversationId = conversationId, - linesOfCodeAccepted = insertedLines, - charactersOfCodeAccepted = insertedCharacters, - ) - } - }.onFailure { /* Noop on diff telemetry failure */ } - - insertNewFiles(newFilePaths) - - applyDeleteFiles(newDeletedFiles) - - ReferenceLogController.addReferenceLog(references, project) - - // Taken from https://intellij-support.jetbrains.com/hc/en-us/community/posts/206118439-Refresh-after-external-changes-to-project-structure-and-sources - VfsUtil.markDirtyAndRefresh(true, true, true, context.addressableRoot) - } - -// Suppressing because insertNewFiles needs to be a suspend function in order to be tested - @Suppress("RedundantSuspendModifier") - suspend fun insertNewFiles( - filePaths: List, - ) { - filePaths.forEach { - resolveAndCreateOrUpdateFile(context.addressableRoot.toNioPath(), it.zipFilePath, it.fileContent) - it.changeApplied = true - } - } - -// Suppressing because applyDeleteFiles needs to be a suspend function in order to be tested - @Suppress("RedundantSuspendModifier") - suspend fun applyDeleteFiles( - deletedFiles: List, - ) { - deletedFiles.forEach { - resolveAndDeleteFile(context.addressableRoot.toNioPath(), it.zipFilePath) - it.changeApplied = true - } - } - - suspend fun disableFileList( - filePaths: List, - deletedFiles: List, - messenger: MessagePublisher, - ) { - if (this._codeResultMessageId.isNullOrEmpty()) { - return - } - - updateFilesPaths(filePaths, deletedFiles, messenger, disableFileActions = true) - this._codeResultMessageId = null - } - - fun sendMetricDataTelemetry(operationName: MetricDataOperationName, result: MetricDataResult) { - featureDevService.sendFeatureDevMetricData(operationName.toString(), result.toString()) - } - - suspend fun send(msg: String): Interaction { - // When the task/"thing to do" hasn't been set yet, we want it to be the incoming message - if (task.isEmpty() && msg.isNotEmpty()) { - task = msg - } - - _latestMessage = msg - return nextInteraction(msg) - } - - private suspend fun nextInteraction(msg: String): Interaction { - var action = - SessionStateAction( - task = task, - msg = msg, - token = sessionState.token, - ) - val resp = sessionState.interact(action) - if (resp.nextState != null) { - // Approach may have been changed after the interaction - val newApproach = sessionState.approach - - // Move to the next state - _state = resp.nextState - - // If approach was changed then we need to set it in the next state and this state - sessionState.approach = newApproach - } - return resp.interaction - } - - private fun getSessionStateConfig(): SessionStateConfig = SessionStateConfig( - conversationId = this.conversationId, - repoContext = this.context, - featureDevService = this.featureDevService, - ) - - val conversationId: String - get() { - if (_conversationId == null) { - throw ConversationIdNotFoundException(operation = "Session", desc = "Conversation ID not found") - } else { - return _conversationId as String - } - } - - val conversationIdUnsafe: String? - get() = _conversationId - - val sessionState: SessionState - get() { - if (_state == null) { - throw Error("State should be initialized before it's read") - } else { - return _state as SessionState - } - } - - var latestMessage: String - get() = this._latestMessage - set(value) { - this._latestMessage = value - } - - val retries: Int - get() = codegenRetries - - fun decreaseRetries() { - codegenRetries -= 1 - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionState.kt deleted file mode 100644 index 936b7a88d3a..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionState.kt +++ /dev/null @@ -1,18 +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.amazonqFeatureDev.session - -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource - -interface SessionState { - val tabID: String - val phase: SessionStatePhase? - var token: CancellationTokenSource? - var codeGenerationRemainingIterationCount: Int? - var codeGenerationTotalIterationCount: Int? - var currentIteration: Int? - var approach: String - var diffMetricsProcessed: DiffMetricsProcessed - suspend fun interact(action: SessionStateAction): SessionStateInteraction -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionStateTypes.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionStateTypes.kt deleted file mode 100644 index c1f8f24c87e..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionStateTypes.kt +++ /dev/null @@ -1,84 +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.amazonqFeatureDev.session - -import com.fasterxml.jackson.annotation.JsonValue -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService -import software.aws.toolkits.jetbrains.services.cwc.messages.RecommendationContentSpan - -data class SessionStateAction( - val task: String, - val msg: String, - val token: CancellationTokenSource? = null, -) - -data class Interaction( - val content: String?, - val interactionSucceeded: Boolean, -) - -data class SessionStateInteraction( - val nextState: T? = null, - val interaction: Interaction, -) - -enum class SessionStatePhase( - @field:JsonValue val json: String, -) { - INIT("Init"), - CODEGEN("Codegen"), -} - -data class SessionStateConfig( - val conversationId: String, - val repoContext: FeatureDevSessionContext, - val featureDevService: FeatureDevService, -) - -data class NewFileZipInfo( - val zipFilePath: String, - val fileContent: String, - var rejected: Boolean, - var changeApplied: Boolean, -) - -data class DeletedFileInfo( - val zipFilePath: String, // The string is the path of the file to be deleted - var rejected: Boolean, - var changeApplied: Boolean, -) - -data class CodeGenerationResult( - var newFiles: List, - var deletedFiles: List, - var references: List, - var codeGenerationRemainingIterationCount: Int? = null, - var codeGenerationTotalIterationCount: Int? = null, -) - -data class CodeReferenceGenerated( - val licenseName: String? = null, - val repository: String? = null, - val url: String? = null, - val recommendationContentSpan: RecommendationContentSpan? = null, -) - -@Suppress("ConstructorParameterNaming") // Unfortunately, this is exactly how the string json is received and is needed for parsing. -data class CodeGenerationStreamResult( - var new_file_contents: Map, - var deleted_files: List, - var references: List, -) - -@Suppress("ConstructorParameterNaming") // Unfortunately, this is exactly how the string json is received and is needed for parsing. -data class ExportTaskAssistResultArchiveStreamResult( - var code_generation_result: CodeGenerationStreamResult, -) - -data class DiffMetricsProcessed( - var accepted: HashSet, - var generated: HashSet, -) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt deleted file mode 100644 index ed0f7eab7be..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt +++ /dev/null @@ -1,33 +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.amazonqFeatureDev.storage - -import com.intellij.openapi.project.Project -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session - -class ChatSessionStorage { - private val sessions = mutableMapOf() - - private fun createSession(tabId: String, project: Project): Session { - val session = Session(tabId, project) - sessions[tabId] = session - return session - } - - @Synchronized fun getSession(tabId: String, project: Project): Session = sessions[tabId] ?: createSession(tabId, project) - - fun deleteSession(tabId: String) { - sessions.remove(tabId) - } - - // Find all sessions that are currently waiting to be authenticated - fun getAuthenticatingSessions(): List = this.sessions.values.filter { it.isAuthenticating } - - fun deleteAllSessions() { - sessions.values.forEach { session -> - session.sessionState.token?.cancel() - } - sessions.clear() - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CancellationTokenSource.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CancellationTokenSource.kt deleted file mode 100644 index ff05dc3c12e..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CancellationTokenSource.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.services.amazonqFeatureDev.util - -class CancellationToken { - private var isCancelled = false - - fun isCancellationRequested(): Boolean = isCancelled - - internal fun cancel() { - isCancelled = true - } -} - -class CancellationTokenSource { - private val _token = CancellationToken() - val token: CancellationToken - get() = _token - - fun cancel() { - _token.cancel() - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CodeReferenceExtensions.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CodeReferenceExtensions.kt deleted file mode 100644 index f2fad9024d2..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CodeReferenceExtensions.kt +++ /dev/null @@ -1,15 +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.amazonqFeatureDev.util - -import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererLicenseInfoManager -import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference - -fun CodeReference.licenseText(): String { - val licenseLink = CodeWhispererLicenseInfoManager.getInstance().getLicenseLink(this.licenseName.orEmpty()) - - return "" + - "${this.licenseName}" + - " license from repository ${this.repository}" -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/DiffMetrics.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/DiffMetrics.kt deleted file mode 100644 index 46b44bf7f90..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/DiffMetrics.kt +++ /dev/null @@ -1,84 +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.amazonqFeatureDev.util - -import com.intellij.diff.comparison.ComparisonManager -import com.intellij.diff.comparison.ComparisonPolicy -import com.intellij.diff.fragments.LineFragment -import com.intellij.openapi.progress.EmptyProgressIndicator -import io.ktor.utils.io.core.toByteArray -import java.nio.charset.Charset -import java.security.MessageDigest - -data class DiffMetrics( - val insertedLines: Int, - val insertedCharacters: Int, -) - -fun lineEnding(content: String, curr: Int, end: Int): Int { - require(curr <= end) { "curr must be within end of range" } - require(end <= content.length) { "end must be within content" } - - return if (curr == end) { - -1 - } else if (content[curr] == '\r') { - if ((curr + 1 < end) && (content[curr + 1] == '\n')) { - 2 - } else { - 1 - } - } else if (content[curr] == '\n') { - 1 - } else { - -1 - } -} - -fun getDiffMetrics(before: String, after: String): DiffMetrics { - val comparisonManager = ComparisonManager.getInstance() - val fragments = comparisonManager.compareLines( - before, - after, - ComparisonPolicy.IGNORE_WHITESPACES, - EmptyProgressIndicator() - ) - - var accLineCount = 0 - var accCharCount = 0 - - fragments.forEach { fragment: LineFragment -> - var curr = fragment.startOffset2 - val end = fragment.endOffset2 - - while (curr < end) { - accLineCount += 1 - - // Consume leading whitespace: - while (curr < end && lineEnding(after, curr, end) == -1 && after[curr].isWhitespace()) curr++ - - // Consume through EOL: - val lineContentStart = curr - while (curr < end && lineEnding(after, curr, end) == -1) curr++ - var lineContentEnd = curr - curr += maxOf(lineEnding(after, curr, end), 0) - - // Walk back trailing whitespace and record character count before continuing to next line: - while (lineContentEnd > lineContentStart && after[lineContentEnd - 1].isWhitespace()) lineContentEnd-- - accCharCount += lineContentEnd - lineContentStart - } - } - - return DiffMetrics( - insertedLines = accLineCount, - insertedCharacters = accCharCount, - ) -} - -fun getChangeIdentifier(filePath: String, before: String, after: String): String { - val hash = MessageDigest.getInstance("SHA-1") - hash.update(filePath.toByteArray(Charset.forName("UTF-8"))) - hash.update(before.toByteArray(Charset.forName("UTF-8"))) - hash.update(after.toByteArray(Charset.forName("UTF-8"))) - return hash.digest().joinToString("") { "%02x".format(it) } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevControllerUtil.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevControllerUtil.kt deleted file mode 100644 index b336d0e085e..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevControllerUtil.kt +++ /dev/null @@ -1,45 +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.amazonqFeatureDev.util - -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpIcons -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpStatusType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpTypes -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase -import software.aws.toolkits.resources.message - -enum class InsertAction { - ALL, - REMAINING, - CONTINUE, - AUTO_CONTINUE, -} - -fun getFollowUpOptions(phase: SessionStatePhase?, type: InsertAction): List { - when (phase) { - SessionStatePhase.CODEGEN -> { - return listOf( - FollowUp( - pillText = when (type) { - InsertAction.ALL -> message("amazonqFeatureDev.follow_up.insert_all_code") - InsertAction.REMAINING -> message("amazonqFeatureDev.follow_up.insert_remaining_code") - InsertAction.CONTINUE -> message("amazonqFeatureDev.follow_up.continue") - InsertAction.AUTO_CONTINUE -> message("amazonqFeatureDev.follow_up.continue") - }, - type = FollowUpTypes.INSERT_CODE, - icon = FollowUpIcons.Ok, - status = FollowUpStatusType.Success - ), - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.provide_feedback_and_regenerate"), - type = FollowUpTypes.PROVIDE_FEEDBACK_AND_REGENERATE_CODE, - icon = FollowUpIcons.Refresh, - status = FollowUpStatusType.Info - ) - ) - } - else -> return emptyList() - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevService.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevService.kt deleted file mode 100644 index 3f6e1f52c8b..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevService.kt +++ /dev/null @@ -1,307 +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.amazonqFeatureDev.util - -import com.fasterxml.jackson.databind.DeserializationFeature -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.SendTelemetryEventResponse -import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationResponse -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.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.amazonqFeatureDev.clients.FeatureDevClient -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeGenerationStreamResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.ExportTaskAssistResultArchiveStreamResult -import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl -import software.aws.toolkits.telemetry.AmazonqTelemetry -import software.aws.toolkits.telemetry.Result - -private val logger = getLogger() - -class FeatureDevService(val proxyClient: FeatureDevClient, val project: Project) { - private val objectMapper = jacksonObjectMapper() - .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false) - - fun createConversation(): String { - val startTime = System.currentTimeMillis() - var failureReason: String? = null - var failureReasonDesc: String? = null - var result: Result = Result.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 = Result.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()}" } - - // BE service will throw ServiceQuota if conversation limit is reached. API Front-end will throw Throttling with this message if conversation limit is reached - if ( - e is software.amazon.awssdk.services.codewhispererruntime.model.ServiceQuotaExceededException || - ( - e is software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException && - e.message?.contains("reached for this month.") == true - ) - ) { - 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): - CreateUploadUrlResponse { - try { - logger.debug { "Executing createUploadUrl with conversationId $conversationId" } - val uploadUrlResponse = proxyClient.createTaskAssistUploadUrl( - conversationId, - contentChecksumSha256, - contentLength, - uploadId - ) - 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) { - 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) - } - } - - fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, message: String, codeGenerationId: String?, currentCodeGenerationId: String?): - StartTaskAssistCodeGenerationResponse { - try { - logger.debug { "Executing startTaskAssistCodeGeneration with conversationId: $conversationId , uploadId: $uploadId" } - val startCodeGenerationResponse = proxyClient.startTaskAssistCodeGeneration( - conversationId, - uploadId, - message, - codeGenerationId, - currentCodeGenerationId ?: "EMPTY_CURRENT_CODE_GENERATION_ID" - ) - - 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) - } - // BE service will throw ServiceQuota if code generation iteration limit is reached - else if (e is software.amazon.awssdk.services.codewhispererruntime.model.ServiceQuotaExceededException || ( - e is software.amazon.awssdk.services.codewhispererruntime.model.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) { - 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): CodeGenerationStreamResult { - 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: ExportTaskAssistResultArchiveStreamResult - try { - val result = exportResponse.reduce { acc, next -> acc + next } // To map the result it is needed to combine the full byte array - parsedResult = objectMapper.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.code_generation_result - } - - fun sendFeatureDevEvent(conversationId: String) { - val sendFeatureDevTelemetryEventResponse: SendTelemetryEventResponse - try { - sendFeatureDevTelemetryEventResponse = proxyClient.sendFeatureDevTelemetryEvent(conversationId) - val requestId = sendFeatureDevTelemetryEventResponse.responseMetadata().requestId() - logger.debug { - "$FEATURE_NAME: succesfully sent feature dev telemetry: ConversationId: $conversationId RequestId: $requestId" - } - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: failed to send feature dev telemetry" } - } - } - - fun sendFeatureDevMetricData(operationName: String, result: String) { - val sendFeatureDevTelemetryEventResponse: SendTelemetryEventResponse - try { - sendFeatureDevTelemetryEventResponse = proxyClient.sendFeatureDevMetricData(operationName, result) - val requestId = sendFeatureDevTelemetryEventResponse.responseMetadata().requestId() - logger.debug { - "$FEATURE_NAME: succesfully sent feature dev metric data: OperationName: $operationName Result: $result RequestId: $requestId" - } - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: failed to send feature dev metric data" } - } - } - - fun sendFeatureDevCodeGenerationEvent(conversationId: String, linesOfCodeGenerated: Int, charactersOfCodeGenerated: Int) { - val sendFeatureDevTelemetryEventResponse: SendTelemetryEventResponse - try { - sendFeatureDevTelemetryEventResponse = proxyClient - .sendFeatureDevCodeGenerationEvent(conversationId, linesOfCodeGenerated, charactersOfCodeGenerated) - val requestId = sendFeatureDevTelemetryEventResponse.responseMetadata().requestId() - logger.debug { - "$FEATURE_NAME: successfully sent feature dev code generation telemetry: ConversationId: $conversationId RequestId: $requestId" - } - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: failed to send feature dev code generation telemetry" } - } - } - - fun sendFeatureDevCodeAcceptanceEvent(conversationId: String, linesOfCodeAccepted: Int, charactersOfCodeAccepted: Int) { - val sendFeatureDevTelemetryEventResponse: SendTelemetryEventResponse - try { - sendFeatureDevTelemetryEventResponse = proxyClient - .sendFeatureDevCodeAcceptanceEvent(conversationId, linesOfCodeAccepted, charactersOfCodeAccepted) - val requestId = sendFeatureDevTelemetryEventResponse.responseMetadata().requestId() - logger.debug { - "$FEATURE_NAME: successfully sent feature dev code acceptance telemetry: ConversationId: $conversationId RequestId: $requestId" - } - } catch (e: Exception) { - logger.warn(e) { "$FEATURE_NAME: failed to send feature dev code acceptance telemetry" } - } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FileUtils.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FileUtils.kt deleted file mode 100644 index c1013708828..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FileUtils.kt +++ /dev/null @@ -1,44 +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.amazonqFeatureDev.util - -import com.intellij.openapi.fileChooser.FileChooser -import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory -import com.intellij.openapi.project.Project -import com.intellij.openapi.vfs.CharsetToolkit -import com.intellij.openapi.vfs.VirtualFile -import java.io.File -import java.nio.charset.Charset -import java.nio.file.Path -import kotlin.io.path.createDirectories -import kotlin.io.path.deleteIfExists -import kotlin.io.path.writeBytes - -/** - * FileUtils.kt - * - * These are utility functions to abstact file IO calls for testing purposes. - */ - -fun resolveAndCreateOrUpdateFile(projectRootPath: Path, relativeFilePath: String, fileContent: String) { - val filePath = projectRootPath.resolve(relativeFilePath) - filePath.parent.createDirectories() // Create directories if needed - filePath.writeBytes(fileContent.toByteArray(Charsets.UTF_8)) -} - -fun resolveAndDeleteFile(projectRootPath: Path, relativePath: String) { - val filePath = projectRootPath.resolve(relativePath) - filePath.deleteIfExists() -} - -fun selectFolder(project: Project, openOn: VirtualFile): VirtualFile? { - val fileChooserDescriptor = FileChooserDescriptorFactory.createSingleFolderDescriptor() - return FileChooser.chooseFile(fileChooserDescriptor, project, openOn) -} - -fun readFileToString(file: File): String { - val charsetToolkit = CharsetToolkit(file.readBytes(), Charset.forName("UTF-8"), false) - val charset = charsetToolkit.guessEncoding(4096) - return file.readText(charset) -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/UploadArtifact.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/UploadArtifact.kt deleted file mode 100644 index cb435409320..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/UploadArtifact.kt +++ /dev/null @@ -1,54 +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.amazonqFeatureDev.util - -import com.intellij.util.io.HttpRequests -import software.amazon.awssdk.utils.IoUtils -import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.core.AwsClientManager -import software.aws.toolkits.jetbrains.services.amazonq.APPLICATION_ZIP -import software.aws.toolkits.jetbrains.services.amazonq.AWS_KMS -import software.aws.toolkits.jetbrains.services.amazonq.CONTENT_SHA256 -import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION -import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevOperation -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.UploadCodeException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.UploadURLExpired -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient -import java.io.File -import java.net.HttpURLConnection - -private val logger = getLogger() -fun uploadArtifactToS3(url: String, fileToUpload: File, checksumSha256: String, contentLength: Long, kmsArn: String?) { - try { - HttpRequests.put(url, APPLICATION_ZIP).userAgent(AwsClientManager.getUserAgent()).tuner { - it.setRequestProperty("Content-Type", APPLICATION_ZIP) - it.setRequestProperty("Content-Length", contentLength.toString()) - it.setRequestProperty(CONTENT_SHA256, checksumSha256) - if (!kmsArn.isNullOrEmpty()) { - it.setRequestProperty(SERVER_SIDE_ENCRYPTION, AWS_KMS) - it.setRequestProperty(SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID, kmsArn) - } - } - .connect { - val connection = it.connection as HttpURLConnection - connection.setFixedLengthStreamingMode(fileToUpload.length()) - IoUtils.copy(fileToUpload.inputStream(), connection.outputStream) - } - } catch (err: HttpRequests.HttpStatusException) { - logger.warn(err) { "$FEATURE_NAME: Failed to upload code to S3" } - when (err.statusCode) { - 403 -> throw UploadURLExpired(operation = FeatureDevOperation.UploadToS3.toString(), desc = "Upload URL expired or forbidden") - else -> throw UploadCodeException(operation = FeatureDevOperation.UploadToS3.toString(), desc = "Failed to upload code to S3") - } - } -} - -fun deleteUploadArtifact(payload: File) { - if (!payload.delete()) { - logger.warn { "Unable to delete upload artifact." } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ReferenceLogController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ReferenceLogController.kt index b46660d0793..70842bd457d 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ReferenceLogController.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ReferenceLogController.kt @@ -7,7 +7,6 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionReference import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionReferencePosition -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager @@ -38,22 +37,4 @@ object ReferenceLogController { ) } } - - fun addReferenceLog(codeReferences: List?, project: Project) { - val manager = CodeWhispererCodeReferenceManager.getInstance(project) - - // TODO flare: hook /dev references with flare correctly, this is only a compile error fix which is not tested - codeReferences?.forEach { reference -> - val cwReferences = InlineCompletionReference( - referenceName = reference.repository.orEmpty(), - referenceUrl = reference.url.orEmpty(), - licenseName = reference.licenseName.orEmpty(), - position = InlineCompletionReferencePosition( - startCharacter = reference.recommendationContentSpan?.start ?: 0, - endCharacter = reference.recommendationContentSpan?.end ?: 0, - ) - ) - manager.addReferenceLogPanelEntry(reference = cwReferences, null, null, null) - } - } } diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevSessionContextTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevSessionContextTest.kt deleted file mode 100644 index 088de821afe..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevSessionContextTest.kt +++ /dev/null @@ -1,141 +0,0 @@ -// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -import com.intellij.testFramework.RuleChain -import org.assertj.core.api.Assertions.assertThat -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.mockito.kotlin.mock -import org.mockito.kotlin.whenever -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService -import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule -import software.aws.toolkits.jetbrains.utils.rules.addFileToModule -import java.util.zip.ZipFile - -data class FileCase(val path: String, val content: String = "", val shouldInclude: Boolean = true) - -class FeatureDevSessionContextTest : FeatureDevTestBase(HeavyJavaCodeInsightTestFixtureRule()) { - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, disposableRule) - private lateinit var featureDevSessionContext: FeatureDevSessionContext - private lateinit var featureDevService: FeatureDevService - - @Before - fun setUp() { - featureDevService = mock() - whenever(featureDevService.project).thenReturn(projectRule.project) - featureDevSessionContext = FeatureDevSessionContext(featureDevService.project, 1024) - } - - private fun fileCases(autoBuildEnabled: Boolean) = listOf( - FileCase(path = ".gitignore"), - FileCase(path = ".gradle/cached.jar", shouldInclude = false), - FileCase(path = "src/MyClass.java"), - FileCase(path = "gradlew"), - FileCase(path = "gradlew.bat"), - FileCase(path = "README.md"), - FileCase(path = "settings.gradle"), - FileCase(path = "build.gradle"), - FileCase(path = "gradle/wrapper/gradle-wrapper.properties"), - FileCase(path = "builder/GetTestBuilder.java"), - FileCase(path = ".aws-sam/build/function1", shouldInclude = false), - FileCase(path = ".gem/specs.rb", shouldInclude = false), - FileCase(path = "archive.zip"), - FileCase(path = "output.bin"), - FileCase(path = "images/logo.png"), - FileCase(path = "assets/header.jpg"), - FileCase(path = "icons/menu.svg"), - FileCase(path = "license.txt"), - FileCase(path = "License.md"), - FileCase(path = "node_modules/express", shouldInclude = false), - FileCase(path = "build/outputs", shouldInclude = false), - FileCase(path = "dist/bundle.js", shouldInclude = false), - FileCase(path = "gradle/wrapper/gradle-wrapper.jar"), - FileCase(path = "big-file.txt", content = "blob".repeat(1024 * 1024), shouldInclude = false), - FileCase(path = "devfile.yaml", shouldInclude = autoBuildEnabled), - ) - - private fun checkZipProject(autoBuildEnabled: Boolean, fileCases: Iterable, onBeforeZip: (() -> Unit)? = null) { - fileCases.forEach { - projectRule.fixture.addFileToModule(module, it.path, it.content) - } - - onBeforeZip?.invoke() - - val zipResult = featureDevSessionContext.getProjectZip(autoBuildEnabled) - val zipPath = zipResult.payload.path - - val zippedFiles = mutableSetOf() - ZipFile(zipPath).use { zipFile -> - for (entry in zipFile.entries()) { - if (!entry.name.endsWith("/")) { - zippedFiles.add(entry.name) - } - } - } - - // The input file paths are relative to the workspaceRoot, however the zip content is relative to the addressableRoot: - val addressableRoot = featureDevSessionContext.addressableRoot.path - val workspaceRoot = featureDevSessionContext.workspaceRoot.path - val base = addressableRoot.removePrefix(workspaceRoot).removePrefix("/") - fun addressablePathOf(path: String) = path.removePrefix(base).removePrefix("/") - - fileCases.forEach { - if (it.shouldInclude) { - assertThat(zippedFiles).contains(addressablePathOf(it.path)) - } else { - assertThat(zippedFiles).doesNotContain(addressablePathOf(it.path)) - } - } - } - - @Test - fun `test zip with autoBuild enabled`() { - checkZipProject(autoBuildEnabled = true, fileCases(autoBuildEnabled = true)) - } - - @Test - fun `test zip with autoBuild disabled`() { - checkZipProject(autoBuildEnabled = false, fileCases(autoBuildEnabled = false)) - } - - @Test - fun `test content is included when selection root is workspace root`() { - val fileCases = listOf( - FileCase(path = "file.txt", shouldInclude = true), - FileCase(path = "project/file.txt", shouldInclude = true), - FileCase(path = "deep/nested/file.txt", shouldInclude = true) - ) - - checkZipProject(autoBuildEnabled = false, fileCases = fileCases, onBeforeZip = { - featureDevSessionContext.selectionRoot = featureDevSessionContext.workspaceRoot - }) - } - - @Test - fun `test content is included within selection root which is deeper than content root`() { - val fileCases = listOf(FileCase(path = "project/module/deep/file.txt", shouldInclude = true)) - - checkZipProject(autoBuildEnabled = false, fileCases = fileCases, onBeforeZip = { - featureDevSessionContext.selectionRoot = featureDevSessionContext.workspaceRoot.findFileByRelativePath("project/module/deep") - ?: error("Failed to find fixture") - }) - } - - @Test - fun `test content is excluded outside of selection root`() { - val fileCases = listOf( - FileCase(path = "project/module/file.txt", shouldInclude = true), - FileCase(path = "project/outside/no.txt", shouldInclude = false), - ) - - checkZipProject(autoBuildEnabled = false, fileCases = fileCases, onBeforeZip = { - featureDevSessionContext.selectionRoot = featureDevSessionContext.workspaceRoot.findFileByRelativePath("project/module") - ?: error("Failed to find fixture") - }) - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt deleted file mode 100644 index ee655405844..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt +++ /dev/null @@ -1,200 +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.amazonqFeatureDev - -import com.intellij.openapi.module.Module -import com.intellij.openapi.project.Project -import com.intellij.openapi.project.modules -import com.intellij.openapi.vfs.VirtualFile -import com.intellij.testFramework.DisposableRule -import com.intellij.testFramework.replaceService -import org.junit.Before -import org.junit.Rule -import org.mockito.Mockito -import org.mockito.kotlin.any -import org.mockito.kotlin.doReturn -import org.mockito.kotlin.mock -import org.mockito.kotlin.spy -import org.mockito.kotlin.whenever -import software.amazon.awssdk.awscore.DefaultAwsResponseMetadata -import software.amazon.awssdk.awscore.util.AwsHeader -import software.amazon.awssdk.services.codewhispererruntime.model.CodeGenerationStatus -import software.amazon.awssdk.services.codewhispererruntime.model.CodeGenerationWorkflowStatus -import software.amazon.awssdk.services.codewhispererruntime.model.CreateTaskAssistConversationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse -import software.amazon.awssdk.services.codewhispererruntime.model.GetTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse -import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationResponse -import software.aws.toolkits.core.TokenConnectionSettings -import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider -import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeGenerationStreamResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated -import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule -import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule -import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule -import software.aws.toolkits.jetbrains.utils.rules.addModule -import java.io.File - -open class FeatureDevTestBase( - @Rule @JvmField - val projectRule: CodeInsightTestFixtureRule = JavaCodeInsightTestFixtureRule(), -) { - @Rule - @JvmField - val disposableRule = DisposableRule() - - internal lateinit var project: Project - internal lateinit var module: Module - internal lateinit var clientAdaptorSpy: FeatureDevClient - internal lateinit var toolkitConnectionManager: ToolkitConnectionManager - - internal val testRequestId = "test_aws_request_id" - internal val testConversationId = "1234" - internal val userMessage = "test-user-message" - internal val codeGenerationId = "1234" - internal val testUploadId = "5678" - internal val testRepositorySize = 20.0 // Picked a random size - internal val otherStatus = "Other" - internal val testTabId = "test-tab-id" - internal val testFilePaths = mapOf(Pair("test.ts", "This is a comment")) - internal val testRunCommandLogPath = ".amazonq/dev/run_command.log" - internal val testLogPath = mapOf(Pair(testRunCommandLogPath, "This is a log")) - internal val testDeletedFiles = listOf("deleted.ts") - internal val testReferences = listOf(CodeReferenceGenerated()) - internal val testChecksumSha = "test-sha" - internal val testContentLength: Long = 40 - - internal val exampleSendTelemetryEventResponse = - SendTelemetryEventResponse - .builder() - .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as SendTelemetryEventResponse - - internal val exampleCreateTaskAssistConversationResponse = - CreateTaskAssistConversationResponse - .builder() - .conversationId(testConversationId) - .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as CreateTaskAssistConversationResponse - - internal val exampleCreateUploadUrlResponse = - CreateUploadUrlResponse - .builder() - .uploadUrl("https://smth.com") - .uploadId(testUploadId) - .kmsKeyArn("0000000000000000000000000000000000:key/1234abcd") - .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as CreateUploadUrlResponse - - internal val exampleStartTaskAssistConversationResponse = - StartTaskAssistCodeGenerationResponse - .builder() - .conversationId(testConversationId) - .codeGenerationId(codeGenerationId) - .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as StartTaskAssistCodeGenerationResponse - - internal val exampleGetTaskAssistConversationResponse = - GetTaskAssistCodeGenerationResponse - .builder() - .conversationId(testConversationId) - .codeGenerationStatus( - CodeGenerationStatus - .builder() - .status(CodeGenerationWorkflowStatus.IN_PROGRESS) - .currentStage("InProgress") - .build(), - ).responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as GetTaskAssistCodeGenerationResponse - - internal val exampleCompleteGetTaskAssistCodeGenerationResponse = - GetTaskAssistCodeGenerationResponse - .builder() - .conversationId(testConversationId) - .codeGenerationStatus( - CodeGenerationStatus - .builder() - .status(CodeGenerationWorkflowStatus.COMPLETE) - .currentStage("Complete") - .build(), - ).codeGenerationRemainingIterationCount(2) - .codeGenerationTotalIterationCount(3) - .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as GetTaskAssistCodeGenerationResponse - - internal val exampleFailedGetTaskAssistCodeGenerationResponse = - GetTaskAssistCodeGenerationResponse - .builder() - .conversationId(testConversationId) - .codeGenerationStatus( - CodeGenerationStatus - .builder() - .status(CodeGenerationWorkflowStatus.FAILED) - .currentStage("Failed") - .build(), - ).responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as GetTaskAssistCodeGenerationResponse - - internal val exampleOtherGetTaskAssistCodeGenerationResponse = - GetTaskAssistCodeGenerationResponse - .builder() - .conversationId(testConversationId) - .codeGenerationStatus( - CodeGenerationStatus - .builder() - .status(CodeGenerationWorkflowStatus.UNKNOWN_TO_SDK_VERSION) - .currentStage(otherStatus) - .build(), - ).responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId))) - .build() as GetTaskAssistCodeGenerationResponse - - internal val exampleExportResultArchiveResponse = mutableListOf(byteArrayOf(100)) - - internal val exampleExportTaskAssistResultArchiveResponse: CodeGenerationStreamResult = CodeGenerationStreamResult(emptyMap(), emptyList(), emptyList()) - - @Before - open fun setup() { - project = projectRule.project - toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project)) - val provider = mock() - val mockBearerProvider = - mock { - doReturn(provider).whenever(it).delegate - } - val connectionSettingsMock = - mock { - whenever(it.tokenProvider).thenReturn(mockBearerProvider) - } - val toolkitConnection = - mock { - doReturn(connectionSettingsMock).whenever(it).getConnectionSettings() - } - doReturn(toolkitConnection).whenever(toolkitConnectionManager).activeConnectionForFeature(any()) - project.replaceService(ToolkitConnectionManager::class.java, toolkitConnectionManager, disposableRule.disposable) - clientAdaptorSpy = spy(FeatureDevClient.getInstance(project)) - project.replaceService(FeatureDevClient::class.java, clientAdaptorSpy, disposableRule.disposable) - - module = project.modules.firstOrNull() ?: if (projectRule is HeavyJavaCodeInsightTestFixtureRule) { - projectRule.fixture.addModule("module1") - } else { - TODO() - } - - val virtualFileMock = Mockito.mock(VirtualFile::class.java) - doReturn("dummy/path").whenever(virtualFileMock).path - } - - companion object { - fun String.toResourceFile(): File { - val uri = - FeatureDevTestBase::class.java.getResource("/amazonqFeatureDev/$this")?.toURI() - ?: throw AssertionError("Unable to locate test resource $this file.") - return File(uri) - } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClientTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClientTest.kt deleted file mode 100644 index 5d789caa0d3..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClientTest.kt +++ /dev/null @@ -1,171 +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.amazonqFeatureDev.clients - -import com.intellij.testFramework.RuleChain -import com.intellij.testFramework.replaceService -import org.assertj.core.api.Assertions.assertThat -import org.junit.After -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.mockito.kotlin.any -import org.mockito.kotlin.argumentCaptor -import org.mockito.kotlin.doReturn -import org.mockito.kotlin.mock -import org.mockito.kotlin.stub -import org.mockito.kotlin.verify -import org.mockito.kotlin.verifyNoInteractions -import org.mockito.kotlin.whenever -import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient -import software.amazon.awssdk.services.codewhispererruntime.model.CreateTaskAssistConversationRequest -import software.amazon.awssdk.services.codewhispererruntime.model.CreateTaskAssistConversationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse -import software.amazon.awssdk.services.codewhispererruntime.model.GetTaskAssistCodeGenerationRequest -import software.amazon.awssdk.services.codewhispererruntime.model.GetTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationRequest -import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationResponse -import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClient -import software.amazon.awssdk.services.codewhispererstreaming.model.ExportResultArchiveRequest -import software.amazon.awssdk.services.codewhispererstreaming.model.ExportResultArchiveResponseHandler -import software.amazon.awssdk.services.codewhispererstreaming.model.GenerateTaskAssistPlanRequest -import software.amazon.awssdk.services.codewhispererstreaming.model.GenerateTaskAssistPlanResponseHandler -import software.amazon.awssdk.services.ssooidc.SsoOidcClient -import software.aws.toolkits.core.TokenConnectionSettings -import software.aws.toolkits.core.utils.test.aString -import software.aws.toolkits.jetbrains.core.MockClientManagerRule -import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection -import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile -import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule -import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.settings.AwsSettings -import java.util.UUID -import java.util.concurrent.CompletableFuture - -class FeatureDevClientTest : FeatureDevTestBase() { - val mockClientManagerRule = MockClientManagerRule() - private val mockCredentialRule = MockCredentialManagerRule() - private val authManagerRule = MockToolkitAuthManagerRule() - - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, mockCredentialRule, mockClientManagerRule, disposableRule) - - private lateinit var bearerClient: CodeWhispererRuntimeClient - private lateinit var streamingBearerClient: CodeWhispererStreamingAsyncClient - private lateinit var amazonQStreamingClient: AmazonQStreamingClient - private lateinit var ssoClient: SsoOidcClient - - private lateinit var featureDevClient: FeatureDevClient - private lateinit var connectionManager: ToolkitConnectionManager - private var isTelemetryEnabledDefault: Boolean = false - - @Before - override fun setup() { - super.setup() - featureDevClient = FeatureDevClient.getInstance(projectRule.project) - ssoClient = mockClientManagerRule.create() - - bearerClient = mockClientManagerRule.create().stub { - on { createTaskAssistConversation(any()) } doReturn exampleCreateTaskAssistConversationResponse - on { createUploadUrl(any()) } doReturn exampleCreateUploadUrlResponse - on { startTaskAssistCodeGeneration(any()) } doReturn exampleStartTaskAssistConversationResponse - on { getTaskAssistCodeGeneration(any()) } doReturn exampleGetTaskAssistConversationResponse - } - - streamingBearerClient = mockClientManagerRule.create().stub { - on { - generateTaskAssistPlan(any(), any()) - } doReturn CompletableFuture.completedFuture(mock()) // void type can't be instantiated - - on { - exportResultArchive(any(), any()) - } doReturn CompletableFuture.completedFuture(mock()) // void type can't be instantiated - } - - amazonQStreamingClient = mock() - projectRule.project.replaceService(AmazonQStreamingClient::class.java, amazonQStreamingClient, disposableRule.disposable) - - val mockConnection = mock() - whenever(mockConnection.getConnectionSettings()) doReturn mock() - - connectionManager = mock { - on { - activeConnectionForFeature(any()) - } doReturn authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), listOf("scopes"))) as AwsBearerTokenConnection - } - projectRule.project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable) - - isTelemetryEnabledDefault = AwsSettings.getInstance().isTelemetryEnabled - } - - @After - fun tearDown() { - AwsSettings.getInstance().isTelemetryEnabled = isTelemetryEnabledDefault - } - - @Test - fun `check createTaskAssistConversation`() { - val actual = featureDevClient.createTaskAssistConversation() - argumentCaptor().apply { - verify(bearerClient).createTaskAssistConversation(capture()) - verifyNoInteractions(streamingBearerClient) - assertThat(actual).isInstanceOf(CreateTaskAssistConversationResponse::class.java) - assertThat(actual).usingRecursiveComparison().comparingOnlyFields("conversationID") - .isEqualTo(exampleCreateTaskAssistConversationResponse) - } - } - - @Test - fun `check createTaskAssistUploadUrl`() { - val testContentLength: Long = 42 - - val actual = featureDevClient.createTaskAssistUploadUrl(testConversationId, "test-sha", testContentLength, "uploadId") - - argumentCaptor().apply { - verify(bearerClient).createUploadUrl(capture()) - verifyNoInteractions(streamingBearerClient) - assertThat(actual).isInstanceOf(CreateUploadUrlResponse::class.java) - assertThat(actual).usingRecursiveComparison().comparingOnlyFields("uploadUrl", "uploadId", "kmsKeyArn") - .isEqualTo(exampleCreateUploadUrlResponse) - } - } - - @Test - fun `check startTaskAssistCodeGeneration`() { - val actual = featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - "test-upload-id", - "test-user-message", - currentCodeGenerationId = UUID.randomUUID() - .toString(), - codeGenerationId = codeGenerationId - ) - - argumentCaptor().apply { - verify(bearerClient).startTaskAssistCodeGeneration(capture()) - verifyNoInteractions(streamingBearerClient) - assertThat(actual).isInstanceOf(StartTaskAssistCodeGenerationResponse::class.java) - assertThat(actual).usingRecursiveComparison().comparingOnlyFields("conversationId", "codeGenerationId") - .isEqualTo(exampleStartTaskAssistConversationResponse) - } - } - - @Test - fun `check getTaskAssistCodeGeneration`() { - val actual = featureDevClient.getTaskAssistCodeGeneration(testConversationId, "test-code-generation-id") - - argumentCaptor().apply { - verify(bearerClient).getTaskAssistCodeGeneration(capture()) - verifyNoInteractions(streamingBearerClient) - assertThat(actual).isInstanceOf(GetTaskAssistCodeGenerationResponse::class.java) - assertThat(actual).usingRecursiveComparison().comparingOnlyFields("conversationId", "codeGenerationStatus") - .isEqualTo(exampleGetTaskAssistConversationResponse) - } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt deleted file mode 100644 index dac4edc2f7a..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt +++ /dev/null @@ -1,927 +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.amazonqFeatureDev.controller - -import com.intellij.testFramework.LightVirtualFile -import com.intellij.testFramework.RuleChain -import com.intellij.testFramework.replaceService -import io.mockk.coVerify -import io.mockk.coVerifyOrder -import io.mockk.every -import io.mockk.just -import io.mockk.mockk -import io.mockk.mockkObject -import io.mockk.mockkStatic -import io.mockk.runs -import io.mockk.unmockkAll -import io.mockk.verify -import kotlinx.coroutines.test.runTest -import org.assertj.core.api.Assertions.assertThat -import org.junit.After -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.junit.jupiter.api.assertThrows -import org.mockito.kotlin.any -import org.mockito.kotlin.doNothing -import org.mockito.kotlin.doReturn -import org.mockito.kotlin.inOrder -import org.mockito.kotlin.mock -import org.mockito.kotlin.reset -import org.mockito.kotlin.spy -import org.mockito.kotlin.times -import org.mockito.kotlin.whenever -import software.aws.toolkits.jetbrains.common.util.selectFolder -import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext -import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController -import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededStates -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeIterationLimitException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ContentLengthException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.EmptyPatchException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.GuardrailsException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MetricDataOperationName -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MetricDataResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MonthlyConversationLimitError -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.NoChangeRequiredException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.PromptRefusalException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ThrottlingException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ZipFileCorruptedException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FeatureDevMessageType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpStatusType -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpTypes -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.IncomingFeatureDevMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswer -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendChatInputEnabledMessage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendCodeResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendSystemPrompt -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateFileComponent -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeGenerationState -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DiffMetricsProcessed -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Interaction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStateConfig -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3 -import software.aws.toolkits.resources.message -import software.aws.toolkits.telemetry.AmazonqTelemetry -import org.mockito.kotlin.verify as mockitoVerify - -class FeatureDevControllerTest : FeatureDevTestBase() { - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, disposableRule) - - private lateinit var controller: FeatureDevController - private lateinit var messenger: MessagePublisher - private lateinit var chatSessionStorage: ChatSessionStorage - private lateinit var appContext: AmazonQAppInitContext - private lateinit var authController: AuthController - private lateinit var spySession: Session - private lateinit var featureDevClient: FeatureDevClient - - private val newFileContents = - listOf( - NewFileZipInfo("test.ts", "This is a comment", false, false), - NewFileZipInfo("test2.ts", "This is a rejected file", true, false), - ) - private val deletedFiles = - listOf( - DeletedFileInfo("delete.ts", false, false), - DeletedFileInfo("delete2.ts", true, false), - ) - - @Before - override fun setup() { - super.setup() - - featureDevClient = mock() - messenger = mock() - chatSessionStorage = mock() - projectRule.project.replaceService(FeatureDevClient::class.java, featureDevClient, disposableRule.disposable) - appContext = - mock { - on { project }.thenReturn(project) - on { messagesFromAppToUi }.thenReturn(messenger) - } - authController = spy(AuthController()) - doReturn(AuthNeededStates()).`when`(authController).getAuthNeededStates(any()) - spySession = spy(Session(testTabId, project)) - - mockkStatic( - MessagePublisher::sendAnswer, - MessagePublisher::sendSystemPrompt, - MessagePublisher::sendUpdatePlaceholder, - MessagePublisher::sendChatInputEnabledMessage, - MessagePublisher::sendCodeResult, - MessagePublisher::updateFileComponent, - ) - - mockkStatic("software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.UploadArtifactKt") - every { uploadArtifactToS3(any(), any(), any(), any(), any()) } just runs - - controller = spy(FeatureDevController(appContext, chatSessionStorage, authController)) - } - - @After - fun clear() { - unmockkAll() - } - - @Test - fun `test new tab opened`() { - val message = IncomingFeatureDevMessage.NewTabCreated("new-tab-created", testTabId) - spySession = spy(Session("tabId", project)) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - reset(authController) // needed to have actual logic to test the isAuthenticating later - - runTest { - controller.processNewTabCreatedMessage(message) - } - mockitoVerify(authController, times(1)).getAuthNeededStates(project) - mockitoVerify(chatSessionStorage, times(1)).getSession(testTabId, project) - assertThat(spySession.isAuthenticating).isTrue() - } - - @Test - fun `test newTask and closeSession followUp`() { - /* - Testing both followups together as they share logic, atm they could be verified together. - */ - val followUp = FollowUp(FollowUpTypes.NEW_TASK, pillText = "Work on new task") - val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command") - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - doNothing().whenever(chatSessionStorage).deleteSession(any()) - - mockkObject(AmazonqTelemetry) - every { AmazonqTelemetry.endChat(amazonqConversationId = any(), amazonqEndOfTheConversationLatency = any()) } just runs - - runTest { - spySession.preloader(messenger) - controller.processFollowupClickedMessage(message) - } - - mockitoVerify(chatSessionStorage, times(1)).deleteSession(testTabId) - - coVerifyOrder { - messenger.sendAnswer(testTabId, message("amazonqFeatureDev.chat_message.ask_for_new_task"), messageType = FeatureDevMessageType.Answer) - messenger.sendUpdatePlaceholder(testTabId, message("amazonqFeatureDev.placeholder.new_plan")) - } - - verify( - exactly = 1, - ) { AmazonqTelemetry.endChat(amazonqConversationId = testConversationId, amazonqEndOfTheConversationLatency = any(), createTime = any()) } - } - - @Test - fun `test provideFeedbackAndRegenerateCode`() = - runTest { - val followUp = FollowUp(FollowUpTypes.PROVIDE_FEEDBACK_AND_REGENERATE_CODE, pillText = "Regenerate code") - val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command") - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - - mockkObject(AmazonqTelemetry) - every { AmazonqTelemetry.isProvideFeedbackForCodeGen(amazonqConversationId = any(), enabled = any()) } just runs - - spySession.preloader(messenger) - controller.processFollowupClickedMessage(message) - - coVerifyOrder { - AmazonqTelemetry.isProvideFeedbackForCodeGen(amazonqConversationId = testConversationId, enabled = true, createTime = any()) - messenger.sendAsyncEventProgress(testTabId, inProgress = false) - messenger.sendAnswer( - tabId = testTabId, - message = message("amazonqFeatureDev.code_generation.provide_code_feedback"), - messageType = FeatureDevMessageType.Answer, - canBeVoted = true, - ) - messenger.sendUpdatePlaceholder(testTabId, message("amazonqFeatureDev.placeholder.provide_code_feedback")) - } - } - - @Test - fun `test insertCode`() = - runTest { - val followUp = FollowUp(FollowUpTypes.INSERT_CODE, pillText = "Insert code") - val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command") - - var featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - mockkObject(AmazonqTelemetry) - every { - AmazonqTelemetry.isAcceptedCodeChanges(amazonqNumberOfFilesAccepted = any(), amazonqConversationId = any(), enabled = any()) - } just runs - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - whenever(spySession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "test-command", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - 0, - 0, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - - doReturn(Unit).whenever(spySession).insertChanges(any(), any(), any()) - doReturn(Unit).whenever(spySession).insertNewFiles(any()) - doReturn(Unit).whenever(spySession).applyDeleteFiles(any()) - - spySession.preloader(messenger) - controller.processFollowupClickedMessage(message) - - mockitoVerify( - spySession, - times(1), - ).insertChanges(newFileContents, deletedFiles, testReferences) // updates for all files - coVerifyOrder { - AmazonqTelemetry.isAcceptedCodeChanges( - amazonqNumberOfFilesAccepted = 2.0, // it should be 2 files per test setup - amazonqConversationId = spySession.conversationId, - enabled = true, - createTime = any(), - ) - - // insert changes for only non rejected files - spySession.insertNewFiles(listOf(newFileContents[0])) - spySession.applyDeleteFiles(listOf(deletedFiles[0])) - - spySession.updateFilesPaths( - filePaths = newFileContents, - deletedFiles = deletedFiles, - messenger - ) - messenger.sendAnswer( - tabId = testTabId, - message = message("amazonqFeatureDev.code_generation.updated_code"), - messageType = FeatureDevMessageType.Answer, - canBeVoted = true, - ) - messenger.sendSystemPrompt( - testTabId, - listOf( - FollowUp(FollowUpTypes.NEW_TASK, message("amazonqFeatureDev.follow_up.new_task"), status = FollowUpStatusType.Info), - FollowUp(FollowUpTypes.CLOSE_SESSION, message("amazonqFeatureDev.follow_up.close_session"), status = FollowUpStatusType.Info), - FollowUp(FollowUpTypes.GENERATE_DEV_FILE, message("amazonqFeatureDev.follow_up.generate_dev_file"), status = FollowUpStatusType.Info) - ), - ) - messenger.sendUpdatePlaceholder(testTabId, message("amazonqFeatureDev.placeholder.additional_improvements")) - } - } - - @Test - fun `test handleChat onCodeGeneration succeeds to create files`() = - runTest { - val mockInteraction = mock() - val featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - mockkObject(AmazonqTelemetry) - val mockSession = mock() - whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) - whenever(mockSession.conversationId).thenReturn(testConversationId) - whenever(mockSession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "test-command", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - coVerifyOrder { - messenger.sendAsyncEventProgress(testTabId, true, message("amazonqFeatureDev.chat_message.start_code_generation_retry")) - messenger.sendAnswer(testTabId, message("amazonqFeatureDev.chat_message.requesting_changes"), messageType = FeatureDevMessageType.AnswerStream) - messenger.sendUpdatePlaceholder(testTabId, message("amazonqFeatureDev.placeholder.generating_code")) - messenger.sendCodeResult(testTabId, testUploadId, newFileContents, deletedFiles, testReferences) - messenger.sendSystemPrompt(testTabId, getFollowUpOptions(SessionStatePhase.CODEGEN, InsertAction.ALL)) - messenger.sendUpdatePlaceholder(testTabId, message("amazonqFeatureDev.placeholder.after_code_generation")) - messenger.sendAsyncEventProgress(testTabId, false) - messenger.sendChatInputEnabledMessage(testTabId, false) - } - } - - @Test(expected = RuntimeException::class) - fun `test handleChat onCodeGeneration throws error when sending message to state`() = - runTest { - val mockSession = mock() - - whenever(mockSession.send(userMessage)).thenThrow(RuntimeException()) - whenever(mockSession.conversationId).thenReturn(testConversationId) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - coVerifyOrder { - messenger.sendAsyncEventProgress(testTabId, true, message("amazonqFeatureDev.chat_message.start_code_generation")) - messenger.sendAnswer(testTabId, message("amazonqFeatureDev.chat_message.requesting_changes"), messageType = FeatureDevMessageType.AnswerStream) - messenger.sendUpdatePlaceholder(testTabId, message("amazonqFeatureDev.placeholder.generating_code")) - messenger.sendAsyncEventProgress(testTabId, false) - messenger.sendChatInputEnabledMessage(testTabId, false) - } - } - - @Test - fun `test handleChat onCodeGeneration doesn't return any files with retries`() = - runTest { - val filePaths = emptyList() - val deletedFiles = emptyList() - - val mockInteraction = mock() - - val mockSession = mock() - whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) - whenever(mockSession.conversationId).thenReturn(testConversationId) - whenever(mockSession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "", - mock(), - filePaths, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - whenever(mockSession.retries).thenReturn(3) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - coVerifyOrder { - messenger.sendAnswer(testTabId, message("amazonqFeatureDev.code_generation.no_file_changes"), messageType = FeatureDevMessageType.Answer) - messenger.sendSystemPrompt( - testTabId, - listOf(FollowUp(FollowUpTypes.RETRY, message("amazonqFeatureDev.follow_up.retry"), status = FollowUpStatusType.Warning)), - ) - messenger.sendChatInputEnabledMessage(testTabId, false) - } - } - - @Test - fun `test handleChat onCodeGeneration doesn't return any files no retries`() = - runTest { - val filePaths = emptyList() - val deletedFiles = emptyList() - - val mockInteraction = mock() - - val mockSession = mock() - whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) - whenever(mockSession.conversationId).thenReturn(testConversationId) - whenever(mockSession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "", - mock(), - filePaths, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - whenever(mockSession.retries).thenReturn(0) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - coVerifyOrder { - messenger.sendAnswer(testTabId, message("amazonqFeatureDev.code_generation.no_file_changes"), messageType = FeatureDevMessageType.Answer) - messenger.sendSystemPrompt(testTabId, emptyList()) - messenger.sendChatInputEnabledMessage(testTabId, false) - } - } - - @Test - fun `test handleChat onCodeGeneration sends correct add code messages`() = runTest { - val totalIterations = 10 - - for (remainingIterations in 0 until totalIterations) { - val message = if (remainingIterations > 2) { - message("amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code_or_feedback") - } else if (remainingIterations > 0) { - message( - "amazonqFeatureDev.code_generation.iteration_counts", - remainingIterations, - totalIterations, - ) - } else { - message( - "amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code", - remainingIterations, - totalIterations, - ) - } - val mockSession = mock() - val featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - val mockInteraction = mock() - whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) - whenever(mockSession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "test-command", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 1, - messenger, - remainingIterations, - totalIterations, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - coVerify { - messenger.sendAnswer( - tabId = testTabId, - messageType = FeatureDevMessageType.Answer, - message = message - ) - } - } - } - - @Test - fun `test handleChat onCodeGeneration sends correct messages after cancellation`() = runTest { - val totalIterations = 10 - - for (remainingIterations in -1 until totalIterations) { - // remainingIterations < 0 is to represent the null case - val message = if (remainingIterations > 2 || remainingIterations < 0) { - message("amazonqFeatureDev.code_generation.stopped_code_generation_no_iteration_count_display") - } else if (remainingIterations > 0) { - message( - "amazonqFeatureDev.code_generation.stopped_code_generation", - remainingIterations, - totalIterations, - ) - } else { - message( - "amazonqFeatureDev.code_generation.stopped_code_generation_no_iterations", - remainingIterations, - totalIterations, - ) - } - val mockSession = mock() - val featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - val mockInteraction = mock() - val token = CancellationTokenSource() - token.cancel() - whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) - whenever(mockSession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - token, - "test-command", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 1, - messenger, - (if (remainingIterations < 0) null else remainingIterations), - totalIterations, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - coVerify { - messenger.sendAnswer( - tabId = testTabId, - messageType = FeatureDevMessageType.Answer, - message = message - ) - } - } - } - - @Test - fun `test handleChat onCodeGeneration sends success metrics`() = runTest { - val mockSession = mock() - val featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - val mockInteraction = mock() - whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) - whenever(mockSession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "test-command", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - - controller.onCodeGeneration(mockSession, userMessage, testTabId) - - val mockInOrder = inOrder(mockSession) - - mockInOrder.verify(mockSession).sendMetricDataTelemetry( - MetricDataOperationName.StartCodeGeneration, - MetricDataResult.Success - - ) - mockInOrder.verify(mockSession).sendMetricDataTelemetry( - MetricDataOperationName.EndCodeGeneration, - MetricDataResult.Success - ) - } - - @Test - fun `test handleChat onCodeGeneration sends correct failure metrics for different errors`() = runTest { - data class ErrorTestCase( - val error: Exception, - val expectedMetricResult: MetricDataResult, - ) - - val testCases = listOf( - ErrorTestCase( - EmptyPatchException("EmptyPatchException", "Empty patch"), - MetricDataResult.LlmFailure - ), - ErrorTestCase( - GuardrailsException(operation = "GenerateCode", desc = "Failed guardrails"), - MetricDataResult.Error - ), - ErrorTestCase( - PromptRefusalException(operation = "GenerateCode", desc = "Prompt refused"), - MetricDataResult.Error - ), - ErrorTestCase( - NoChangeRequiredException(operation = "GenerateCode", desc = "No changes needed"), - MetricDataResult.Error - ), - ErrorTestCase( - ThrottlingException(operation = "GenerateCode", desc = "Request throttled"), - MetricDataResult.Error - ), - ErrorTestCase( - MonthlyConversationLimitError(message = "Monthly limit reached", operation = "GenerateCode", desc = "Monthly limit reached"), - MetricDataResult.Error - ), - ErrorTestCase( - CodeIterationLimitException(operation = "GenerateCode", desc = "Code iteration limit reached"), - MetricDataResult.Error - ), - ErrorTestCase( - ContentLengthException(operation = "GenerateCode", desc = "Repo size is exceeding the limits"), - MetricDataResult.Error - ), - ErrorTestCase( - ZipFileCorruptedException(operation = "GenerateCode", desc = "Zipped file is corrupted"), - MetricDataResult.Fault - ), - ErrorTestCase( - FeatureDevException(message = "Resource not found", operation = "GenerateCode", desc = null), - MetricDataResult.Fault - ), - ErrorTestCase( - RuntimeException("Unknown error"), - MetricDataResult.Fault - ) - ) - - testCases.forEach { (error, expectedResult) -> - val mockSession = mock() - whenever(mockSession.send(userMessage)).thenThrow(error) - whenever(mockSession.sessionState).thenReturn( - CodeGenerationState( - testTabId, - "", - mock(), - testUploadId, - 0, - 0.0, - messenger, - token = CancellationTokenSource(), - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()) - ) - ) - - assertThrows { - controller.onCodeGeneration(mockSession, userMessage, testTabId) - } - - val mockInOrder = inOrder(mockSession) - - mockInOrder.verify(mockSession).sendMetricDataTelemetry( - MetricDataOperationName.StartCodeGeneration, - MetricDataResult.Success - - ) - mockInOrder.verify(mockSession).sendMetricDataTelemetry( - MetricDataOperationName.EndCodeGeneration, - expectedResult - ) - } - } - - @Test - fun `test processFileClicked handles file rejection`() = - runTest { - val message = IncomingFeatureDevMessage.FileClicked(testTabId, newFileContents[0].zipFilePath, "", "reject-change") - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - whenever(spySession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "", - mock(), - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - - controller.processFileClicked(message) - - val newFileContentsCopy = newFileContents.toMutableList() - newFileContentsCopy[0] = newFileContentsCopy[0].copy() - newFileContentsCopy[0].rejected = true - newFileContentsCopy[0].changeApplied = false - coVerify { messenger.updateFileComponent(testTabId, newFileContentsCopy, deletedFiles, "") } - } - - @Test - fun `test processFileClicked handles file acceptance`() = - runTest { - val featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - whenever(spySession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - doReturn(testConversationId).`when`(spySession).conversationId - doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any()) - - mockkObject(AmazonqTelemetry) - every { - AmazonqTelemetry.isAcceptedCodeChanges( - amazonqNumberOfFilesAccepted = 1.0, - amazonqConversationId = testConversationId, - enabled = true, - credentialStartUrl = any() - ) - } just runs - - // Accept first file: - controller.processFileClicked(IncomingFeatureDevMessage.FileClicked(testTabId, newFileContents[0].zipFilePath, "", "accept-change")) - - val newFileContentsCopy = newFileContents.toList() - newFileContentsCopy[0].rejected = false - newFileContentsCopy[0].changeApplied = true - coVerify { messenger.updateFileComponent(testTabId, newFileContents, deletedFiles, "") } - - mockitoVerify( - spySession, - times(1), - ).insertChanges(listOf(newFileContents[0]), listOf(), testReferences) - - // Does not continue automatically, because files are remaining: - mockitoVerify( - controller, - times(0), - ).insertCode(testTabId) - } - - @Test - fun `test processFileClicked automatically continues when last file is accepted`() = - runTest { - val featureDevService = mockk() - val repoContext = mock() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - whenever(spySession.sessionState).thenReturn( - PrepareCodeGenerationState( - testTabId, - CancellationTokenSource(), - "", - sessionStateConfig, - newFileContents, - deletedFiles, - testReferences, - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ), - ) - doReturn(testConversationId).`when`(spySession).conversationId - doReturn(Unit).`when`(spySession).insertChanges(any(), any(), any()) - - mockkObject(AmazonqTelemetry) - every { - AmazonqTelemetry.isAcceptedCodeChanges( - amazonqNumberOfFilesAccepted = 1.0, - amazonqConversationId = testConversationId, - enabled = true, - credentialStartUrl = any() - ) - } just runs - - val newFileContentsCopy = newFileContents.toList() - newFileContentsCopy[0].rejected = false - newFileContentsCopy[0].changeApplied = true - newFileContentsCopy[1].rejected = false - newFileContentsCopy[1].changeApplied = true - deletedFiles[0].rejected = false - deletedFiles[0].changeApplied = true - - // This is simulating the file already being an accepted state, and accept-change being called redundantly. This is necessary because of the test - // setup, which should be fixed to avoid heavy-handed mocking of the session state (so that we can see the session state be incrementally updated). - deletedFiles[1].rejected = false - deletedFiles[1].changeApplied = true - - // When the last file is accepted: - controller.processFileClicked(IncomingFeatureDevMessage.FileClicked(testTabId, deletedFiles[1].zipFilePath, "", "accept-change")) - - // We auto-continue to the next step with a noop insertCode call: - mockitoVerify( - controller, - times(1), - ).insertCode(testTabId) - } - - @Test - fun `test modifyDefaultSourceFolder customer does not select a folder`() = - runTest { - val followUp = FollowUp(FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, pillText = "Modify default source folder") - val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command") - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - - mockkStatic("software.aws.toolkits.jetbrains.common.util.FileUtilsKt") - every { selectFolder(any(), any()) } returns null - - spySession.preloader(messenger) - controller.processFollowupClickedMessage(message) - - coVerifyOrder { - messenger.sendSystemPrompt( - tabId = testTabId, - followUp = - listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"), - type = FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, - status = FollowUpStatusType.Info, - ), - ), - ) - } - } - - @Test - fun `test modifyDefaultSourceFolder customer selects a folder outside the workspace`() = - runTest { - val followUp = FollowUp(FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, pillText = "Modify default source folder") - val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command") - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - - mockkStatic("software.aws.toolkits.jetbrains.common.util.FileUtilsKt") - every { selectFolder(any(), any()) } returns LightVirtualFile("/path") - - spySession.preloader(messenger) - controller.processFollowupClickedMessage(message) - - coVerifyOrder { - messenger.sendAnswer( - tabId = testTabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.follow_up.incorrect_source_folder"), - ) - messenger.sendSystemPrompt( - tabId = testTabId, - followUp = - listOf( - FollowUp( - pillText = message("amazonqFeatureDev.follow_up.modify_source_folder"), - type = FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, - status = FollowUpStatusType.Info, - ), - ), - ) - } - } - - @Test - fun `test modifyDefaultSourceFolder customer selects a correct sub folder`() = - runTest { - val followUp = FollowUp(FollowUpTypes.MODIFY_DEFAULT_SOURCE_FOLDER, pillText = "Modify default source folder") - val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command") - - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse) - whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession) - - val folder = LightVirtualFile("${spySession.context.workspaceRoot.path.removePrefix("/")}/path/to/sub/folder") - mockkStatic("software.aws.toolkits.jetbrains.common.util.FileUtilsKt") - every { selectFolder(any(), any()) } returns folder - - spySession.preloader(messenger) - controller.processFollowupClickedMessage(message) - - coVerify { - messenger.sendAnswer( - tabId = testTabId, - messageType = FeatureDevMessageType.Answer, - message = message("amazonqFeatureDev.follow_up.modified_source_folder", folder.path), - canBeVoted = true, - ) - } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt deleted file mode 100644 index 8e60630ac21..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt +++ /dev/null @@ -1,199 +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.amazonqFeatureDev.session - -import com.intellij.testFramework.RuleChain -import io.mockk.coEvery -import io.mockk.coVerify -import io.mockk.every -import io.mockk.just -import io.mockk.mockk -import io.mockk.mockkStatic -import io.mockk.runs -import io.mockk.unmockkAll -import io.mockk.verify -import kotlinx.coroutines.test.runTest -import org.assertj.core.api.Assertions.assertThat -import org.assertj.core.api.Assertions.assertThatThrownBy -import org.junit.After -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.mockito.kotlin.mock -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService -import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference -import software.aws.toolkits.resources.message - -class CodeGenerationStateTest : FeatureDevTestBase() { - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, disposableRule) - - private lateinit var codeGenerationState: CodeGenerationState - private lateinit var messenger: MessagePublisher - private val action = SessionStateAction("test-task", userMessage) - private lateinit var featureDevService: FeatureDevService - private lateinit var repoContext: FeatureDevSessionContext - - @Before - override fun setup() { - featureDevService = mockk() - every { featureDevService.project } returns projectRule.project - messenger = mock() - repoContext = mockk() - val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - - codeGenerationState = - CodeGenerationState( - testTabId, - "", - sessionStateConfig, - testUploadId, - 0, - testRepositorySize, - messenger, - token = null, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), - ) - - mockkStatic(MessagePublisher::sendAnswerPart) - coEvery { messenger.sendAnswerPart(any(), any()) } just runs - } - - @After - fun clear() { - unmockkAll() - } - - @Test - fun `test code generated is complete`() { - val action = SessionStateAction("test-task", userMessage) - every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleCompleteGetTaskAssistCodeGenerationResponse - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - coEvery { featureDevService.exportTaskAssistArchiveResult(any()) } returns - CodeGenerationStreamResult(testFilePaths, testDeletedFiles, testReferences) - - runTest { - val actual = codeGenerationState.interact(action) - assertThat(actual.nextState).isInstanceOf(PrepareCodeGenerationState::class.java) - val nextState = actual.nextState as PrepareCodeGenerationState - assertThat(nextState.phase).isEqualTo(SessionStatePhase.CODEGEN) - assertThat(nextState.filePaths).isEqualTo( - listOf(NewFileZipInfo("test.ts", "This is a comment", rejected = false, changeApplied = false)), - ) - assertThat(nextState.deletedFiles).isEqualTo( - listOf(DeletedFileInfo("deleted.ts", rejected = false, changeApplied = false)), - ) - assertThat(nextState.references).isEqualTo(testReferences) - assertThat(nextState.codeGenerationRemainingIterationCount).isEqualTo(2) - assertThat(nextState.codeGenerationTotalIterationCount).isEqualTo(3) - assertThat(actual.interaction.interactionSucceeded).isEqualTo(true) - assertThat(actual.interaction.content).isEqualTo("") - } - assertThat(codeGenerationState.phase).isEqualTo(SessionStatePhase.CODEGEN) - coVerify(exactly = 1) { messenger.sendAnswerPart(testTabId, message("amazonqFeatureDev.code_generation.generating_code")) } - - verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) } - coVerify(exactly = 1) { featureDevService.exportTaskAssistArchiveResult(testConversationId) } - } - - @Test - fun `test generateCode excludes run_command log file`() { - val runCommandLogFileName = "run_command.log" - - val archiveFiles = mapOf( - runCommandLogFileName to "newLog", - "other.ts" to "other content" - ) - val deletedFiles = emptyList() - val references = emptyList() - - every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleCompleteGetTaskAssistCodeGenerationResponse - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - coEvery { featureDevService.exportTaskAssistArchiveResult(any()) } returns - CodeGenerationStreamResult(archiveFiles, listOf("deleted.ts"), listOf(CodeReferenceGenerated())) - - runTest { - val actual = codeGenerationState.interact(action) - val nextState = actual.nextState as PrepareCodeGenerationState - assertThat(nextState.filePaths).contains( - NewFileZipInfo(runCommandLogFileName, "newLog", rejected = false, changeApplied = false) - ) - assertThat(nextState.filePaths).contains( - NewFileZipInfo("other.ts", "other content", rejected = false, changeApplied = false) - ) - } - } - - @Test(expected = FeatureDevException::class) - fun `test code generation failed`() = - runTest { - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleFailedGetTaskAssistCodeGenerationResponse - - codeGenerationState.interact(action) - - verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) } - coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) } - } - - @Test - fun `test code generation returns any other handled status`() { - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleOtherGetTaskAssistCodeGenerationResponse - - assertThatThrownBy { - runTest { - codeGenerationState.interact(action) - } - }.isExactlyInstanceOf(IllegalStateException::class.java).withFailMessage("Unknown status: $otherStatus") - - verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) } - coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) } - } - - @Test - fun `test code generation returns in progress at least once`() { - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - every { - featureDevService.getTaskAssistCodeGeneration(any(), any()) - } returnsMany listOf(exampleGetTaskAssistConversationResponse, exampleCompleteGetTaskAssistCodeGenerationResponse) - coEvery { featureDevService.exportTaskAssistArchiveResult(any()) } returns CodeGenerationStreamResult(emptyMap(), emptyList(), emptyList()) - - runTest { - codeGenerationState.interact(action) - } - - verify(exactly = 2) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) } - coVerify(exactly = 1) { featureDevService.exportTaskAssistArchiveResult(testConversationId) } - } - - @Test - fun `test using all polling count`() { - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleGetTaskAssistConversationResponse - - runTest { - val actual = codeGenerationState.interact(action) - assertThat(actual.nextState).isInstanceOf(PrepareCodeGenerationState::class.java) - val nextState = actual.nextState as PrepareCodeGenerationState - assertThat(nextState.filePaths).isEqualTo(emptyList()) - assertThat(nextState.deletedFiles).isEqualTo(emptyList()) - assertThat(nextState.references).isEqualTo(emptyList()) - assertThat(nextState.codeGenerationRemainingIterationCount).isEqualTo(null) - assertThat(nextState.codeGenerationTotalIterationCount).isEqualTo(null) - assertThat(actual.interaction.interactionSucceeded).isEqualTo(true) - assertThat(actual.interaction.content).isEqualTo("") - } - - verify(exactly = 360) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) } - coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(testConversationId) } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt deleted file mode 100644 index 3df897e053d..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt +++ /dev/null @@ -1,102 +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.amazonqFeatureDev.session - -import com.intellij.testFramework.RuleChain -import io.mockk.coEvery -import io.mockk.every -import io.mockk.just -import io.mockk.mockk -import io.mockk.mockkStatic -import io.mockk.runs -import io.mockk.unmockkAll -import kotlinx.coroutines.test.runTest -import org.assertj.core.api.Assertions.assertThat -import org.junit.After -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.mockito.kotlin.mock -import org.mockito.kotlin.times -import org.mockito.kotlin.verify -import org.mockito.kotlin.whenever -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonq.project.FeatureDevSessionContext -import software.aws.toolkits.jetbrains.services.amazonq.project.ZipCreationResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.deleteUploadArtifact -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3 -import java.io.File - -class PrepareCodeGenerationStateTest : FeatureDevTestBase() { - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, disposableRule) - - private lateinit var prepareCodeGenerationState: PrepareCodeGenerationState - private lateinit var repoContext: FeatureDevSessionContext - private lateinit var sessionStateConfig: SessionStateConfig - private lateinit var messenger: MessagePublisher - private lateinit var featureDevService: FeatureDevService - - @Before - override fun setup() { - repoContext = mock() - featureDevService = mockk() - every { featureDevService.project } returns projectRule.project - messenger = mock() - sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) - - prepareCodeGenerationState = PrepareCodeGenerationState( - "", - CancellationTokenSource(), - "test-approach", - sessionStateConfig, - emptyList(), - emptyList(), - emptyList(), - testUploadId, - 0, - messenger, - diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()) - ) - - mockkStatic("software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.UploadArtifactKt") - every { uploadArtifactToS3(any(), any(), any(), any(), any()) } just runs - every { deleteUploadArtifact(any()) } just runs - - every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleCompleteGetTaskAssistCodeGenerationResponse - every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse - coEvery { featureDevService.exportTaskAssistArchiveResult(any()) } returns exampleExportTaskAssistResultArchiveResponse - - mockkStatic(MessagePublisher::sendAnswerPart) - coEvery { messenger.sendAnswerPart(any(), any()) } just runs - } - - @After - fun clear() { - unmockkAll() - } - - @Test - fun `test interact`() { - val mockFile: File = mock() - val repoZipResult = ZipCreationResult(mockFile, testChecksumSha, testContentLength) - val action = SessionStateAction("test-task", userMessage) - - whenever(repoContext.workspaceRoot).thenReturn(mock()) - whenever(repoContext.getProjectZip(false)).thenReturn(repoZipResult) - every { featureDevService.createUploadUrl(any(), any(), any(), any()) } returns exampleCreateUploadUrlResponse - - runTest { - val actual = prepareCodeGenerationState.interact(action) - assertThat(actual.nextState).isInstanceOf(PrepareCodeGenerationState::class.java) - } - assertThat(prepareCodeGenerationState.phase).isEqualTo(SessionStatePhase.CODEGEN) - verify(repoContext, times(1)).getProjectZip(false) - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionTest.kt deleted file mode 100644 index 2bf1eb7765d..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionTest.kt +++ /dev/null @@ -1,95 +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.amazonqFeatureDev.session - -import com.intellij.openapi.vfs.VfsUtil -import com.intellij.openapi.vfs.VirtualFile -import com.intellij.testFramework.RuleChain -import com.intellij.testFramework.replaceService -import io.mockk.every -import io.mockk.just -import io.mockk.mockkObject -import io.mockk.mockkStatic -import io.mockk.runs -import io.mockk.verify -import kotlinx.coroutines.runBlocking -import kotlinx.coroutines.test.runTest -import org.assertj.core.api.Assertions.assertThat -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.mockito.kotlin.mock -import org.mockito.kotlin.times -import org.mockito.kotlin.verify -import org.mockito.kotlin.whenever -import software.aws.toolkits.jetbrains.common.util.resolveAndCreateOrUpdateFile -import software.aws.toolkits.jetbrains.common.util.resolveAndDeleteFile -import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient -import software.aws.toolkits.jetbrains.services.cwc.controller.ReferenceLogController -import kotlin.io.path.Path - -class SessionTest : FeatureDevTestBase() { - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, disposableRule) - - private lateinit var featureDevClient: FeatureDevClient - private lateinit var session: Session - private lateinit var messenger: MessagePublisher - - @Before - override fun setup() { - featureDevClient = mock() - projectRule.project.replaceService(FeatureDevClient::class.java, featureDevClient, disposableRule.disposable) - session = Session("tabId", projectRule.project) - messenger = mock() - } - - @Test - fun `test session before preloader`() { - assertThat(session.sessionState).isInstanceOf(ConversationNotStartedState::class.java) - assertThat(session.isAuthenticating).isFalse() - } - - @Test - fun `test preloader`() = runTest { - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - - session.preloader(messenger) - assertThat(session.conversationId).isEqualTo(testConversationId) - assertThat(session.sessionState).isInstanceOf(PrepareCodeGenerationState::class.java) - verify(featureDevClient, times(1)).createTaskAssistConversation() - } - - @Test - fun `test insertChanges`() { - mockkStatic("com.intellij.openapi.vfs.VfsUtil") - every { VfsUtil.markDirtyAndRefresh(true, true, true, any()) } just runs - - mockkObject(ReferenceLogController) - every { ReferenceLogController.addReferenceLog(any(), any()) } just runs - - mockkStatic("software.aws.toolkits.jetbrains.common.util.FileUtilsKt") - every { resolveAndDeleteFile(any(), any()) } just runs - every { resolveAndCreateOrUpdateFile(any(), any(), any()) } just runs - - val mockNewFile = listOf(NewFileZipInfo("test.ts", "testContent", rejected = false, changeApplied = false)) - val mockDeletedFile = listOf(DeletedFileInfo("deletedTest.ts", rejected = false, changeApplied = false)) - - val addressableRootPath = Path("src") - session.context.addressableRoot = mock() - whenever(session.context.addressableRoot.toNioPath()).thenReturn(addressableRootPath) - - runBlocking { - session.insertChanges(mockNewFile, mockDeletedFile, emptyList()) - } - - verify(exactly = 1) { resolveAndDeleteFile(addressableRootPath, "deletedTest.ts") } - verify(exactly = 1) { resolveAndCreateOrUpdateFile(addressableRootPath, "test.ts", "testContent") } - verify(exactly = 1) { ReferenceLogController.addReferenceLog(emptyList(), any()) } - verify(exactly = 1) { VfsUtil.markDirtyAndRefresh(true, true, true, session.context.addressableRoot) } - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorageTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorageTest.kt deleted file mode 100644 index d0c0289028e..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorageTest.kt +++ /dev/null @@ -1,44 +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.amazonqFeatureDev.storage - -import org.assertj.core.api.Assertions.assertThat -import org.junit.Before -import org.junit.Test -import org.mockito.kotlin.mock -import org.mockito.kotlin.whenever -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session - -class ChatSessionStorageTest : FeatureDevTestBase() { - - private lateinit var chatSessionStorage: ChatSessionStorage - private lateinit var mockSession: Session - - @Before - override fun setup() { - super.setup() - chatSessionStorage = ChatSessionStorage() - mockSession = mock() - } - - @Test - fun `check getSession for NewSession`() { - val testSession = chatSessionStorage.getSession("tabId", project) - assertThat(testSession).isNotNull() - assertThat(testSession.tabID).isEqualTo("tabId") - assertThat(testSession.project).isEqualTo(project) - } - - @Test - fun `check getSession for ExistingSession`() { - whenever(mockSession.tabID).thenReturn("tab1") - whenever(mockSession.project).thenReturn(projectRule.project) - - val expectedSession = chatSessionStorage.getSession(mockSession.tabID, mockSession.project) - val actualSession = chatSessionStorage.getSession("tab1", project) - assertThat(actualSession).isNotNull() - assertThat(actualSession).isEqualTo(expectedSession) - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/DiffMetricsTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/DiffMetricsTest.kt deleted file mode 100644 index 8b4b5721ba4..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/DiffMetricsTest.kt +++ /dev/null @@ -1,135 +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.amazonqFeatureDev.util - -import com.intellij.testFramework.LightPlatformTestCase -import com.intellij.testFramework.TestApplicationManager - -class DiffMetricsTest : LightPlatformTestCase() { - override fun setUp() { - super.setUp() - TestApplicationManager.getInstance() - } - - fun `test empty input`() { - val metrics = getDiffMetrics("", "") - assertEquals(0, metrics.insertedLines) - assertEquals(0, metrics.insertedCharacters) - } - - fun `test insertions are counted`() { - val before = """ - line1 - line2 - """.trimIndent() - - val after = """ - line1 - inserted - line2 - """.trimIndent() - - val metrics = getDiffMetrics(before, after) - assertEquals(1, metrics.insertedLines) - assertEquals(8, metrics.insertedCharacters) - } - - fun `test modifications are counted`() { - val before = """ - line1 - line2 - line3 - """.trimIndent() - - val after = """ - line1 - modified - line3 - """.trimIndent() - - val metrics = getDiffMetrics(before, after) - assertEquals(1, metrics.insertedLines) - assertEquals(8, metrics.insertedCharacters) - } - - fun `test deletions are counted`() { - val before = """ - line1 - line2 - line3 - """.trimIndent() - - val after = """ - line1 - line3 - """.trimIndent() - - val metrics = getDiffMetrics(before, after) - assertEquals(0, metrics.insertedLines) - assertEquals(0, metrics.insertedCharacters) - } - - fun `test multiline and multiple hunks are counted`() { - val before = """ - line1 - line2 - line3 - """.trimIndent() - - val after = """ - inserted1 - line1 - inserted2 - inserted3 - line3 - inserted4 - """.trimIndent() - - val metrics = getDiffMetrics(before, after) - assertEquals(4, metrics.insertedLines) - assertEquals(36, metrics.insertedCharacters) - } - - fun `test empty lines are counted`() { - val before = "line1" - val after = "line1\n\nline2" - val metrics = getDiffMetrics(before, after) - assertEquals(2, metrics.insertedLines) - assertEquals(5, metrics.insertedCharacters) - } - - fun `test trailing newline is not counted`() { - val before = "line1" - val after = "line1\nline2\n" - val metrics = getDiffMetrics(before, after) - assertEquals(1, metrics.insertedLines) - assertEquals(5, metrics.insertedCharacters) - } - - fun `test newline sequences are counted`() { - val before = "line1" - val after = "line1\nline2\rline3\r\nline4" - val metrics = getDiffMetrics(before, after) - assertEquals(3, metrics.insertedLines) - assertEquals(15, metrics.insertedCharacters) - } - - fun `test leading and trailing whitespace are not counted as characters`() { - val before = "line1\nline2" - val after = "line1\n after " - - val metrics = getDiffMetrics(before, after) - assertEquals(1, metrics.insertedLines) - assertEquals(5, metrics.insertedCharacters) - } - - fun `test ignore whitespace change when performing diff`() { - val before = "line1\nline2" - val after = "line1\n line2" - - val metrics = getDiffMetrics(before, after) - assertEquals(0, metrics.insertedLines) - assertEquals(0, metrics.insertedCharacters) - } -} diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevServiceTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevServiceTest.kt deleted file mode 100644 index 2cccde2f496..00000000000 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevServiceTest.kt +++ /dev/null @@ -1,383 +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.amazonqFeatureDev.util - -import com.intellij.testFramework.RuleChain -import com.intellij.testFramework.replaceService -import kotlinx.coroutines.test.runTest -import org.assertj.core.api.Assertions.assertThat -import org.assertj.core.api.Assertions.assertThatThrownBy -import org.junit.Before -import org.junit.Rule -import org.junit.Test -import org.mockito.kotlin.mock -import org.mockito.kotlin.whenever -import software.amazon.awssdk.awscore.exception.AwsErrorDetails -import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException -import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse -import software.amazon.awssdk.services.codewhispererruntime.model.ServiceQuotaExceededException -import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException -import software.amazon.awssdk.services.codewhispererruntime.model.ValidationException -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.FeatureDevException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MonthlyConversationLimitError -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ServiceException -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeGenerationStreamResult -import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated -import software.aws.toolkits.resources.message - -class FeatureDevServiceTest : FeatureDevTestBase() { - @Rule - @JvmField - val ruleChain = RuleChain(projectRule, disposableRule) - - private lateinit var featureDevClient: FeatureDevClient - private lateinit var featureDevService: FeatureDevService - - private val cwExceptionMsg = "CWException" - private val otherExceptionMsg = "otherException" - - @Before - override fun setup() { - featureDevClient = mock() - projectRule.project.replaceService(FeatureDevClient::class.java, featureDevClient, disposableRule.disposable) - featureDevService = FeatureDevService(featureDevClient, projectRule.project) - } - - @Test - fun `test createConversation`() { - whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse) - val actual = featureDevService.createConversation() - assertThat(actual).isEqualTo(exampleCreateTaskAssistConversationResponse.conversationId()) - } - - @Test - fun `test createConversation with error`() { - whenever(featureDevClient.createTaskAssistConversation()).thenThrow(RuntimeException()) - assertThatThrownBy { - featureDevService.createConversation() - }.isInstanceOf(FeatureDevException::class.java) - } - - @Test - fun `test createUploadUrl`() { - whenever( - featureDevClient.createTaskAssistUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId = testUploadId), - ).thenReturn(exampleCreateUploadUrlResponse) - - val actual = featureDevService.createUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId = testUploadId) - assertThat(actual).isInstanceOf(CreateUploadUrlResponse::class.java) - assertThat(actual) - .usingRecursiveComparison() - .comparingOnlyFields("uploadUrl", "uploadId", "kmsKeyArn") - .isEqualTo(exampleCreateUploadUrlResponse) - } - - @Test - fun `test createUploadUrl with error`() { - whenever( - featureDevClient.createTaskAssistUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId = testUploadId), - ).thenThrow(RuntimeException()) - assertThatThrownBy { - featureDevService.createUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId = testUploadId) - }.isInstanceOf(FeatureDevException::class.java) - } - - @Test - fun `test createUploadUrl with validation error`() { - whenever( - featureDevClient.createTaskAssistUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId = testUploadId), - ).thenThrow( - ValidationException - .builder() - .requestId(testRequestId) - .message("Invalid contentLength") - .awsErrorDetails(AwsErrorDetails.builder().errorMessage("Invalid contentLength").build()) - .build(), - ) - - assertThatThrownBy { - featureDevService.createUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId = testUploadId) - }.isInstanceOf(ContentLengthException::class.java).hasMessage(message("amazonqFeatureDev.content_length.error_text")) - } - - @Test - fun `test getTaskAssistCodeGeneration success`() { - whenever( - featureDevClient.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()), - ).thenReturn(exampleGetTaskAssistConversationResponse) - - val actual = featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) - - assertThat(actual).isEqualTo(exampleGetTaskAssistConversationResponse) - } - - @Test - fun `test getTaskAssistCodeGeneration throws a CodeWhispererRuntimeException`() { - val exampleCWException = - CodeWhispererRuntimeException - .builder() - .awsErrorDetails( - AwsErrorDetails.builder().errorMessage(cwExceptionMsg).build(), - ).build() - whenever(featureDevClient.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString())).thenThrow(exampleCWException) - - assertThatThrownBy { - featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) - }.isExactlyInstanceOf(ServiceException::class.java).withFailMessage(cwExceptionMsg) - } - - @Test - fun `test getTaskAssistCodeGeneration throws another Exception`() { - val exampleOtherException = RuntimeException(otherExceptionMsg) - whenever(featureDevClient.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString())).thenThrow(exampleOtherException) - - assertThatThrownBy { - featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) - }.isExactlyInstanceOf(ServiceException::class.java).withFailMessage(otherExceptionMsg) - } - - @Test - fun `test startTaskAssistConversation success`() { - whenever( - featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ), - ).thenReturn(exampleStartTaskAssistConversationResponse) - - val actual = - featureDevService.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ) - - assertThat(actual).isEqualTo(exampleStartTaskAssistConversationResponse) - } - - @Test - fun `test startTaskAssistConversation throws ThrottlingException`() { - val exampleCWException = - ThrottlingException - .builder() - .awsErrorDetails( - AwsErrorDetails.builder().errorMessage("limit for number of iterations on a code generation").build(), - ).build() - whenever( - featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ), - ).thenThrow(exampleCWException) - - assertThatThrownBy { - featureDevService.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ) - }.isExactlyInstanceOf(CodeIterationLimitException::class.java).withFailMessage( - message("amazonqFeatureDev.code_generation.iteration_limit.error_text"), - ) - } - - @Test - fun `test startTaskAssistConversation throws ThrottlingException, different case`() { - val exampleCWException = - ThrottlingException - .builder() - .awsErrorDetails( - AwsErrorDetails.builder().errorMessage("StartTaskAssistCodeGeneration reached for this month.").build(), - ).build() - whenever( - featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ), - ).thenThrow(exampleCWException) - - assertThatThrownBy { - featureDevService.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ) - }.isExactlyInstanceOf(MonthlyConversationLimitError::class.java).withFailMessage( - message("amazonqFeatureDev.exception.monthly_limit_error"), - ) - } - - @Test - fun `test startTaskAssistConversation throws ServiceQuotaExceededException`() { - val exampleCWException = - ServiceQuotaExceededException - .builder() - .awsErrorDetails( - AwsErrorDetails.builder().errorMessage("service quota exceeded").build(), - ).build() - whenever( - featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ), - ).thenThrow(exampleCWException) - - assertThatThrownBy { - featureDevService.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ) - }.isExactlyInstanceOf(CodeIterationLimitException::class.java).withFailMessage( - message("amazonqFeatureDev.code_generation.iteration_limit.error_text"), - ) - } - - @Test - fun `test startTaskAssistConversation throws another CodeWhispererRuntimeException`() { - val exampleCWException = - CodeWhispererRuntimeException - .builder() - .awsErrorDetails( - AwsErrorDetails.builder().errorMessage(cwExceptionMsg).build(), - ).build() - whenever( - featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ), - ).thenThrow(exampleCWException) - - assertThatThrownBy { - featureDevService.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ) - }.isExactlyInstanceOf(ServiceException::class.java).withFailMessage(cwExceptionMsg) - } - - @Test - fun `test startTaskAssistConversation throws any other exception`() { - val exampleOtherException = RuntimeException(otherExceptionMsg) - whenever( - featureDevClient.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ), - ).thenThrow(exampleOtherException) - - assertThatThrownBy { - featureDevService.startTaskAssistCodeGeneration( - testConversationId, - testUploadId, - userMessage, - codeGenerationId = codeGenerationId, - currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID", - ) - }.isExactlyInstanceOf(ServiceException::class.java).withFailMessage(otherExceptionMsg) - } - - @Test - fun `test exportTaskAssistArchiveResult throws CodeWhispererStreamingException`() { - val exampleCWException = - CodeWhispererRuntimeException - .builder() - .awsErrorDetails( - AwsErrorDetails.builder().errorMessage(cwExceptionMsg).build(), - ).build() - - assertThatThrownBy { - runTest { - whenever(featureDevClient.exportTaskAssistResultArchive(testConversationId)).thenThrow(exampleCWException) - featureDevService.exportTaskAssistArchiveResult(testConversationId) - } - }.isExactlyInstanceOf(ServiceException::class.java).withFailMessage(cwExceptionMsg) - } - - @Test - fun `test exportTaskAssistArchiveResult throws other exception`() { - val exampleOtherException = RuntimeException(otherExceptionMsg) - - assertThatThrownBy { - runTest { - whenever(featureDevClient.exportTaskAssistResultArchive(testConversationId)).thenThrow(exampleOtherException) - featureDevService.exportTaskAssistArchiveResult(testConversationId) - } - }.isExactlyInstanceOf(ServiceException::class.java).withFailMessage(otherExceptionMsg) - } - - @Test - fun `test exportTaskAssistArchiveResult throws an exportParseError`() { - assertThatThrownBy { - runTest { - whenever(featureDevClient.exportTaskAssistResultArchive(testConversationId)).thenReturn(mutableListOf(byteArrayOf(0, 1, 2))) - featureDevService.exportTaskAssistArchiveResult(testConversationId) - } - }.isExactlyInstanceOf(ExportParseException::class.java) - } - - @Test - fun `test exportTaskAssistArchiveResult ignore extra fields in the response`() { - runTest { - val codeGenerationJson = "{" + - "\"code_generation_result\":{\"new_file_contents\":{\"test.ts\":\"contents\"},\"deleted_files\":[],\"references\":[],\"extra_filed\":[]}" + - "}" - - whenever(featureDevClient.exportTaskAssistResultArchive(testConversationId)).thenReturn(mutableListOf(codeGenerationJson.toByteArray())) - - val actual = featureDevService.exportTaskAssistArchiveResult(testConversationId) - assertThat(actual).isInstanceOf(CodeGenerationStreamResult::class.java) - assertThat(actual.new_file_contents).isEqualTo(mapOf(Pair("test.ts", "contents"))) - assertThat(actual.deleted_files).isEqualTo(emptyList()) - assertThat(actual.references).isEqualTo(emptyList()) - } - } - - @Test - fun `test exportTaskAssistArchiveResult returns correct parsed result`() = - runTest { - val codeGenerationJson = "{\"code_generation_result\":{\"new_file_contents\":{\"test.ts\":\"contents\"},\"deleted_files\":[],\"references\":[]}}" - whenever(featureDevClient.exportTaskAssistResultArchive(testConversationId)).thenReturn(mutableListOf(codeGenerationJson.toByteArray())) - - val actual = featureDevService.exportTaskAssistArchiveResult(testConversationId) - assertThat(actual).isInstanceOf(CodeGenerationStreamResult::class.java) - assertThat(actual.new_file_contents).isEqualTo(mapOf(Pair("test.ts", "contents"))) - assertThat(actual.deleted_files).isEqualTo(emptyList()) - assertThat(actual.references).isEqualTo(emptyList()) - } -}