fix(dev): messaging

Author: tverney

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerExtensions.kt

@@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.controller
 
 import com.intellij.notification.NotificationAction
 import org.gradle.tooling.GradleConnector
+import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CODE_GENERATION_RETRY_LIMIT
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FeatureDevMessageType
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp
@@ -21,6 +22,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Delete
 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.SessionState
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.resources.message
@@ -44,19 +46,26 @@ suspend fun FeatureDevController.onCodeGeneration(session: Session, message: Str
             message = message("amazonqFeatureDev.chat_message.requesting_changes"),
             messageType = FeatureDevMessageType.AnswerStream,
         )
+        val state = session.sessionState
+
+        var remainingIterations: Int? = state.codeGenerationRemainingIterationCount
+        var totalIterations: Int? = state.codeGenerationTotalIterationCount
+
+
+        if (state.token?.token()?.isCancellationRequested == true) {
+            this.disposeToken(state, messenger, tabId, remainingIterations, totalIterations)
+            return
+        }
 
         messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.generating_code"))
 
         session.send(message) // Trigger code generation
 
-        val state = session.sessionState
 
         var filePaths: List<NewFileZipInfo> = emptyList()
         var deletedFiles: List<DeletedFileInfo> = emptyList()
         var references: List<CodeReferenceGenerated> = emptyList()
         var uploadId = ""
-        var remainingIterations: Int? = null
-        var totalIterations: Int? = null
 
         when (state) {
             is PrepareCodeGenerationState -> {
@@ -70,17 +79,7 @@ suspend fun FeatureDevController.onCodeGeneration(session: Session, message: Str
         }
 
         if (state.token?.token()?.isCancellationRequested == true) {
-            messenger.sendAnswer(
-                tabId = tabId,
-                messageType = FeatureDevMessageType.Answer,
-                message = message("amazonqFeatureDev.code_generation.stopped_code_generation")
-            )
-            messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true)
-
-            messenger.sendUpdatePlaceholder(
-                tabId = tabId,
-                newPlaceholder = message("amazonqFeatureDev.placeholder.new_plan")
-            )
+            disposeToken(state, messenger, tabId, remainingIterations, totalIterations)
             return
         }
 
@@ -149,6 +148,32 @@ suspend fun FeatureDevController.onCodeGeneration(session: Session, message: Str
     }
 }
 
+private suspend fun FeatureDevController.disposeToken(state: SessionState, messenger: MessagePublisher, tabId: String, remainingIterations: Number?, totalIterations: Number?) {
+    if (state.codeGenerationRemainingIterationCount !== null) {
+        messenger.sendAnswer(
+            tabId = tabId,
+            messageType = FeatureDevMessageType.Answer,
+            message = message("amazonqFeatureDev.code_generation.stopped_code_generation", remainingIterations ?: state.currentIteration as Any,
+                totalIterations ?: CODE_GENERATION_RETRY_LIMIT
+            )
+        )
+    } else {
+        messenger.sendAnswer(
+            tabId = tabId,
+            messageType = FeatureDevMessageType.Answer,
+            message = message("amazonqFeatureDev.code_generation.stopped_code_generation_without_total", state.currentIteration as Any)
+        )
+    }
+
+    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")
 ) {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt

@@ -33,11 +33,11 @@ class CodeGenerationState(
     override var approach: String,
     val config: SessionStateConfig,
     val uploadId: String,
-    val currentIteration: Int,
+    override var currentIteration: Int? = 0,
     val repositorySize: Double,
     val messenger: MessagePublisher,
-    var codeGenerationRemainingIterationCount: Int? = null,
-    var codeGenerationTotalIterationCount: Int? = null,
+    override var codeGenerationRemainingIterationCount: Int? = null,
+    override var codeGenerationTotalIterationCount: Int? = null,
     var currentCodeGenerationId:  UUID? = null,
     override var token: CancellationTokenSource?
     ) : SessionState {
@@ -88,7 +88,7 @@ class CodeGenerationState(
                 filePaths = codeGenerationResult.newFiles,
                 deletedFiles = codeGenerationResult.deletedFiles,
                 references = codeGenerationResult.references,
-                currentIteration = currentIteration + 1,
+                currentIteration = currentIteration?.plus(1),
                 uploadId = uploadId,
                 messenger = messenger,
                 codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount,
@@ -114,20 +114,22 @@ class CodeGenerationState(
 
             throw e
         } finally {
-            AmazonqTelemetry.codeGenerationInvoke(
-                amazonqConversationId = config.conversationId,
-                amazonqCodeGenerationResult = codeGenerationWorkflowStatus.toString(),
-                amazonqGenerateCodeIteration = currentIteration.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)
-            )
+            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)
+                )
+            }
         }
     }
 }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/ConversationNotStartedState.kt

@@ -9,6 +9,9 @@ class ConversationNotStartedState(
     override var approach: String,
     override val tabID: String,
     override var token: CancellationTokenSource?,
+    override var codeGenerationRemainingIterationCount: Int?,
+    override var codeGenerationTotalIterationCount: Int?,
+    override var currentIteration: Int?,
     ) : SessionState {
     override val phase = SessionStatePhase.INIT
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt

@@ -30,10 +30,10 @@ class PrepareCodeGenerationState(
     val deletedFiles: List<DeletedFileInfo>,
     val references: List<CodeReferenceGenerated>,
     var uploadId: String,
-    private val currentIteration: Int,
+    override var currentIteration: Int?,
     private var messenger: MessagePublisher,
-    var codeGenerationRemainingIterationCount: Int? = null,
-    var codeGenerationTotalIterationCount: Int? = null,
+    override var codeGenerationRemainingIterationCount: Int? = null,
+    override var codeGenerationTotalIterationCount: Int? = null,
 ) : SessionState {
     override val phase = SessionStatePhase.CODEGEN
     override suspend fun interact(action: SessionStateAction): SessionStateInteraction {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt

@@ -42,7 +42,7 @@ class Session(val tabID: String, val project: Project) {
         context = FeatureDevSessionContext(project, MAX_PROJECT_SIZE_BYTES)
         proxyClient = FeatureDevClient.getInstance(project)
         featureDevService = FeatureDevService(proxyClient, project)
-        _state = ConversationNotStartedState("", tabID, null)
+        _state = ConversationNotStartedState("", tabID, null, 0, CODE_GENERATION_RETRY_LIMIT, 0)
         isAuthenticating = false
         codegenRetries = CODE_GENERATION_RETRY_LIMIT
     }
@@ -79,7 +79,7 @@ class Session(val tabID: String, val project: Project) {
             filePaths = emptyList(),
             deletedFiles = emptyList(),
             references = emptyList(),
-            currentIteration = 0, // first code gen iteration
+            currentIteration = 1, // first code gen iteration
             uploadId = "", // There is no code gen uploadId so far
             messenger = messenger,
             token = GradleConnector.newCancellationTokenSource()

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionState.kt

@@ -9,6 +9,9 @@ interface SessionState {
     val tabID: String
     val phase: SessionStatePhase?
     var token: CancellationTokenSource?
+    var codeGenerationRemainingIterationCount: Int?
+    var codeGenerationTotalIterationCount: Int?
+    var currentIteration: Int?
     var approach: String
     suspend fun interact(action: SessionStateAction): SessionStateInteraction
 }

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt

@@ -10,13 +10,15 @@ 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.gradle.tooling.CancellationTokenSource
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -29,6 +31,7 @@ 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.services.amazonq.FeatureDevSessionContext
 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
@@ -47,13 +50,16 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendC
 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.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.FeatureDevService
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.selectFolder
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3
@@ -86,6 +92,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
     @Before
     override fun setup() {
         super.setup()
+
         featureDevClient = mock()
         messenger = mock()
         chatSessionStorage = mock()
@@ -199,6 +206,9 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         val followUp = FollowUp(FollowUpTypes.INSERT_CODE, pillText = "Insert code")
         val message = IncomingFeatureDevMessage.FollowupClicked(followUp, testTabId, "", "test-command")
 
+        var featureDevService = mockk<FeatureDevService>()
+        val repoContext = mock<FeatureDevSessionContext>()
+        val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService)
         mockkObject(AmazonqTelemetry)
         every {
             AmazonqTelemetry.isAcceptedCodeChanges(amazonqNumberOfFilesAccepted = any(), amazonqConversationId = any(), enabled = any())
@@ -209,7 +219,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
         whenever(spySession.sessionState).thenReturn(
             PrepareCodeGenerationState(
-                testTabId, "", mock(), newFileContents, deletedFiles, testReferences, testUploadId, 0, messenger
+                testTabId, null, mock(), sessionStateConfig, newFileContents, deletedFiles, testReferences, testUploadId, 0, messenger
             )
         )
         doNothing().`when`(spySession).insertChanges(any(), any(), any())
@@ -248,13 +258,16 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
     @Test
     fun `test handleChat onCodeGeneration succeeds to create files`() = runTest {
         val mockInteraction = mock<Interaction>()
-
+        var featureDevService = mockk<FeatureDevService>()
+        val repoContext = mock<FeatureDevSessionContext>()
+        val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService)
+        mockkObject(AmazonqTelemetry)
         val mockSession = mock<Session>()
         whenever(mockSession.send(userMessage)).thenReturn(mockInteraction)
         whenever(mockSession.conversationId).thenReturn(testConversationId)
         whenever(mockSession.sessionState).thenReturn(
             PrepareCodeGenerationState(
-                testTabId, "", mock(), newFileContents, deletedFiles, testReferences, testUploadId, 0, messenger
+                testTabId, null, "test-command", sessionStateConfig, newFileContents, deletedFiles, testReferences, testUploadId, 0, messenger,
             )
         )
 
@@ -302,7 +315,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         whenever(mockSession.conversationId).thenReturn(testConversationId)
         whenever(mockSession.sessionState).thenReturn(
             PrepareCodeGenerationState(
-                testTabId, "", mock(), filePaths, deletedFiles, testReferences, testUploadId, 0, messenger
+                testTabId, null, "", mock(), filePaths, deletedFiles, testReferences, testUploadId, 0, messenger
             )
         )
         whenever(mockSession.retries).thenReturn(3)
@@ -331,7 +344,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         whenever(mockSession.conversationId).thenReturn(testConversationId)
         whenever(mockSession.sessionState).thenReturn(
             PrepareCodeGenerationState(
-                testTabId, "", mock(), filePaths, deletedFiles, testReferences, testUploadId, 0, messenger
+                testTabId, null, "", mock(), filePaths, deletedFiles, testReferences, testUploadId, 0, messenger
             )
         )
         whenever(mockSession.retries).thenReturn(0)
@@ -353,7 +366,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
         whenever(spySession.sessionState).thenReturn(
             PrepareCodeGenerationState(
-                testTabId, "", mock(), newFileContents, deletedFiles, testReferences, testUploadId, 0, messenger
+                testTabId, null, "", mock(), newFileContents, deletedFiles, testReferences, testUploadId, 0, messenger
             )
         )
 

plugins/amazonq/mynah-ui/src/mynah-ui/ui/texts/constants.ts

@@ -19,7 +19,7 @@ export const uiComponentsTexts = {
     save: 'Save',
     cancel: 'Cancel',
     submit: 'Submit',
-    stopGenerating: 'Stop generating',
+    stopGenerating: 'Stop',
     copyToClipboard: 'Copied to clipboard',
     noMoreTabsTooltip: 'You can only open ten conversation tabs at a time.',
     codeSuggestionWithReferenceTitle: 'Some suggestions contain code with references.',

plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -57,8 +57,9 @@ amazonqFeatureDev.code_generation.notification_message=Your code suggestions fro
 amazonqFeatureDev.code_generation.notification_open_link=Open chat
 amazonqFeatureDev.code_generation.notification_title=Amazon Q Developer Agent for software development
 amazonqFeatureDev.code_generation.provide_code_feedback=How can I improve the code for your use case?
-amazonqFeatureDev.code_generation.stopping_code_generation=Stopping code generation ...
-amazonqFeatureDev.code_generation.stopped_code_generation=You've stopped code generation.
+amazonqFeatureDev.code_generation.stopping_code_generation=Stopping code generation...
+amazonqFeatureDev.code_generation.stopped_code_generation=I stopped generating your code. If you want to continue working on this task, provide another description. You have {0} out of {1} code generations left.
+amazonqFeatureDev.code_generation.stopped_code_generation_without_total=I stopped generating your code. If you want to continue working on this task, provide another description. You have started {0} code generations.
 amazonqFeatureDev.code_generation.updated_code=Okay, I updated your code files. Would you like to work on another task?
 amazonqFeatureDev.content_length.error_text=The folder you selected is too large for me to use as context. Please choose a smaller folder to work on. For more information on quotas, see the [Amazon Q Developer documentation](https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/software-dev.html#quotas). 
 amazonqFeatureDev.error_text=Sorry, we encountered a problem when processing your request.
@@ -87,7 +88,7 @@ amazonqFeatureDev.follow_up.provide_feedback_and_regenerate=Provide feedback & r
 amazonqFeatureDev.follow_up.retry=Retry
 amazonqFeatureDev.follow_up.send_feedback=Send feedback
 amazonqFeatureDev.no_retries.error_text=I'm sorry, I'm having technical difficulties and can't continue at the moment. Please try again later, and share feedback to help me improve.
-amazonqFeatureDev.placeholder.additional_improvements=Choose an option to proceed
+amazonqFeatureDev.placeholder.additional_improvements=Describe your task or issue in detail
 amazonqFeatureDev.placeholder.after_code_generation=Choose an option to proceed
 amazonqFeatureDev.placeholder.after_monthly_limit=Chat input is disabled
 amazonqFeatureDev.placeholder.closed_session=Open a new chat tab to continue