feat(amazonq): Adding UTG build, execute logic for internal users

packages/core/src/amazonq/webview/ui/main.ts

@@ -669,7 +669,7 @@ export const createMynahUI = (
                 mynahUI.addChatItem(tabID, {
                     type: ChatItemType.ANSWER_STREAM,
                 })
-            } else if (tabType === 'gumby') {
+            } else if (tabType === 'gumby' || tabType === 'testgen') {
                 connector.requestAnswer(tabID, {
                     chatMessage: prompt.prompt ?? '',
                 })

packages/core/src/amazonqTest/chat/controller/messenger/messengerUtils.ts

@@ -12,6 +12,7 @@ export enum ButtonActions {
     VIEW_DIFF = 'View-Diff',
     STOP_TEST_GEN = 'Stop-Test-Generation',
     STOP_BUILD = 'Stop-Build-Process',
+    STOP_FIXING_TEST = 'Stop-Fixing-Test',
     PROVIDE_FEEDBACK = 'Provide-Feedback',
 }
 

packages/core/src/amazonqTest/chat/session/session.ts

@@ -3,8 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ShortAnswer, Reference } from '../../../codewhisperer/models/model'
-import { TargetFileInfo, TestGenerationJob } from '../../../codewhisperer/client/codewhispereruserclient'
+import { Reference } from '../../../codewhisperer/models/model'
+import { PackageInfo, TargetFileInfo, TestGenerationJob } from '../../../codewhisperer/client/codewhispereruserclient'
 
 export enum ConversationState {
     IDLE,
@@ -37,13 +37,13 @@ export class Session {
     // Start Test generation
     public isSupportedLanguage: boolean = false
     public conversationState: ConversationState = ConversationState.IDLE
-    public shortAnswer: ShortAnswer | undefined
     public sourceFilePath: string = ''
     public generatedFilePath: string = ''
     public projectRootPath: string = ''
     public fileLanguage: string | undefined = 'plaintext'
     public stopIteration: boolean = false
     public targetFileInfo: TargetFileInfo | undefined
+    public packageInfo: PackageInfo | undefined
     public jobSummary: string = ''
 
     // Telemetry

packages/core/src/amazonqTest/models/constants.ts

@@ -13,15 +13,19 @@ export const testChat = 'testChat'
 
 export const maxUserPromptLength = 4096 // user prompt character limit from MPS and API model.
 
+const baseProgressField: Partial<ProgressField> = {
+    status: 'default',
+    value: -1,
+}
+
 export const cancelTestGenButton: ChatItemButton = {
     id: ButtonActions.STOP_TEST_GEN,
     text: 'Cancel',
     icon: 'cancel' as MynahIcons,
 }
 
 export const testGenProgressField: ProgressField = {
-    status: 'default',
-    value: -1,
+    ...baseProgressField,
     text: 'Generating unit tests...',
     actions: [cancelTestGenButton],
 }
@@ -46,18 +50,29 @@ export const cancelBuildProgressButton: ChatItemButton = {
     icon: 'cancel' as MynahIcons,
 }
 
+export const cancelFixingTestButton: ChatItemButton = {
+    id: ButtonActions.STOP_FIXING_TEST,
+    text: 'Cancel',
+    icon: 'cancel' as MynahIcons,
+}
+
 export const buildProgressField: ProgressField = {
-    status: 'default',
-    value: -1,
-    text: 'Executing...',
+    ...baseProgressField,
+    text: 'Project compiling...',
     actions: [cancelBuildProgressButton],
 }
 
+export const fixingTestProgressField: ProgressField = {
+    ...baseProgressField,
+    text: 'Fixing test failures...',
+    actions: [cancelFixingTestButton],
+}
+
 export const errorProgressField: ProgressField = {
     status: 'error',
     text: 'Error...Input needed',
     value: -1,
-    actions: [cancelBuildProgressButton],
+    actions: [cancelTestGenButton],
 }
 
 export const testGenSummaryMessage = (
@@ -70,67 +85,99 @@ ${planSummary ? `\n\n${planSummary}` : ''}
 `
 
 const checkIcons = {
-    wait: '&#9744;',
-    current: '&#9744;',
+    wait: '&#9203;',
+    current: '&#9203;',
     done: '<span style="color: green;">&#10004;</span>',
     error: '&#10060;',
 }
-
+// TODO: Commenting out this code to do a better UX in the V2 version after science support
+/*
 interface StepStatus {
     step: TestGenerationBuildStep
     status: 'wait' | 'current' | 'done' | 'error'
 }
 
 const stepStatuses: StepStatus[] = []
-
+*/
 export const testGenBuildProgressMessage = (currentStep: TestGenerationBuildStep, status?: string) => {
     const session = ChatSessionManager.Instance.getSession()
-    const statusText = BuildStatus[session.buildStatus].toLowerCase()
-    const icon = session.buildStatus === BuildStatus.SUCCESS ? checkIcons['done'] : checkIcons['error']
-    let message = `Sure. This may take a few minutes and I'll share updates on my progress here.
+    let message = `Sure. This may take a few minutes and I'll update the progress here.\n
 **Progress summary**\n\n`
 
     if (currentStep === TestGenerationBuildStep.START_STEP) {
         return message.trim()
     }
 
+    if (currentStep === TestGenerationBuildStep.RUN_BUILD) {
+        message += `${checkIcons['wait']} Project compiling\n\n`
+    }
+    if (currentStep === TestGenerationBuildStep.FIXING_TEST_CASES && session.buildStatus === BuildStatus.FAILURE) {
+        message += `${checkIcons['wait']} Fixing test failures\n\n`
+    } else if (
+        currentStep >= TestGenerationBuildStep.FIXING_TEST_CASES &&
+        session.buildStatus === BuildStatus.FAILURE
+    ) {
+        message += `${checkIcons['done']} Fixed test failures\n\n`
+    }
+    if (currentStep > TestGenerationBuildStep.RUN_BUILD && session.buildStatus === BuildStatus.SUCCESS) {
+        message += `${checkIcons['done']} Project compiled\n${checkIcons['done']} All tests passed\n\n`
+    }
+    /*
     updateStepStatuses(currentStep, status)
 
     if (currentStep >= TestGenerationBuildStep.RUN_BUILD) {
-        message += `${getIconForStep(TestGenerationBuildStep.RUN_BUILD)} Started build execution\n`
+        message += `${getIconForStep(TestGenerationBuildStep.RUN_BUILD)} ${
+            currentStep === TestGenerationBuildStep.RUN_BUILD
+                ? 'Project compiling\n'
+                : session.buildStatus === BuildStatus.FAILURE
+                  ? 'Unable to compile project\n'
+                  : 'Project compiled\n'
+        }`
     }
 
-    if (currentStep >= TestGenerationBuildStep.RUN_EXECUTION_TESTS) {
-        message += `${getIconForStep(TestGenerationBuildStep.RUN_EXECUTION_TESTS)} Executing tests\n`
+    if (currentStep === TestGenerationBuildStep.RUN_EXECUTION_TESTS) {
+        message += `${getIconForStep(TestGenerationBuildStep.RUN_EXECUTION_TESTS)} Running tests\n`
+    } else if (currentStep >= TestGenerationBuildStep.RUN_EXECUTION_TESTS) {
+        message += `${getIconForStep(TestGenerationBuildStep.RUN_EXECUTION_TESTS)} ${
+            session.buildStatus === BuildStatus.FAILURE ? 'Tests failed\n' : 'Tests passed\n'
+        }`
     }
 
-    if (currentStep >= TestGenerationBuildStep.FIXING_TEST_CASES && session.buildStatus === BuildStatus.FAILURE) {
-        message += `${getIconForStep(TestGenerationBuildStep.FIXING_TEST_CASES)} Fixing errors in tests\n\n`
+    if (currentStep === TestGenerationBuildStep.FIXING_TEST_CASES && session.buildStatus === BuildStatus.FAILURE) {
+        message += `${getIconForStep(TestGenerationBuildStep.FIXING_TEST_CASES)} Fixing test failures\n\n`
+    } else if (
+        currentStep >= TestGenerationBuildStep.FIXING_TEST_CASES &&
+        session.buildStatus === BuildStatus.FAILURE
+    ) {
+        message += `${checkIcons['done']} Fixed test failures\n\n`
     }
-
-    if (currentStep > TestGenerationBuildStep.PROCESS_TEST_RESULTS) {
-        message += `**Test case summary**
-${session.shortAnswer?.testCoverage ? `- Unit test coverage ${session.shortAnswer?.testCoverage}%` : ``}
-${icon} Build ${statusText}
-${icon} Assertion ${statusText}`
-        // TODO: Update Assertion %
+*/
+    if (currentStep > TestGenerationBuildStep.PROCESS_TEST_RESULTS && session.buildStatus === BuildStatus.FAILURE) {
+        message += `**Results**\n
+Amazon Q executed the tests and identified at least one failure. Below are the suggested fixes.\n\n`
     }
 
     return message.trim()
 }
-// TODO: Work on UX to show the build error in the progress message
+/*
 const updateStepStatuses = (currentStep: TestGenerationBuildStep, status?: string) => {
+    const session = ChatSessionManager.Instance.getSession()
     for (let step = TestGenerationBuildStep.INSTALL_DEPENDENCIES; step <= currentStep; step++) {
         const stepStatus: StepStatus = {
             step: step,
             status: 'wait',
         }
 
-        if (step === currentStep) {
-            stepStatus.status = status === 'failed' ? 'error' : 'current'
-        } else if (step < currentStep) {
-            stepStatus.status = 'done'
-        }
+        stepStatus.status =
+            step === currentStep
+                ? status === 'error' || status === 'done'
+                    ? status
+                    : 'current'
+                : step < currentStep
+                  ? session.buildStatus === BuildStatus.FAILURE
+                      ? 'error'
+                      : 'done'
+                  : stepStatus.status
 
         const existingIndex = stepStatuses.findIndex((s) => s.step === step)
         if (existingIndex !== -1) {
@@ -145,3 +192,4 @@ const getIconForStep = (step: TestGenerationBuildStep) => {
     const stepStatus = stepStatuses.find((s) => s.step === step)
     return stepStatus ? checkIcons[stepStatus.status] : checkIcons.wait
 }
+*/

packages/core/src/codewhisperer/commands/startTestGeneration.ts

@@ -44,6 +44,7 @@ export async function startTestGenerationProcess(
             logger.verbose(`Tab ID mismatch: ${tabID} !== ${session.tabID}`)
             return
         }
+        testGenState.setToNotStarted()
         /**
          * Step 1: Zip the project
          */
@@ -53,7 +54,6 @@ export async function startTestGenerationProcess(
             const projectPath = zipUtil.getProjectPath(filePath) ?? ''
             const relativeTargetPath = path.relative(projectPath, filePath)
             session.listOfTestGenerationJobId = []
-            session.shortAnswer = undefined
             session.sourceFilePath = relativeTargetPath
             session.projectRootPath = projectPath
             session.listOfTestGenerationJobId = []
@@ -212,7 +212,7 @@ function runLocalBuild(
         if (spawnResult.stderr) {
             spawnResult.stderr.on('data', async (data) => {
                 const output = data.toString().trim()
-                getLogger().warn(`BUILD ERROR: ${output}`)
+                getLogger().warn(`${output}`)
                 buildLogs += output
             })
         }

packages/core/src/codewhisperer/models/model.ts

@@ -1174,18 +1174,3 @@ export interface Reference {
         end: number
     }
 }
-
-// TODO: remove ShortAnswer because it will be deprecated
-export interface ShortAnswer {
-    testFilePath: string
-    buildCommands: string[]
-    planSummary: string
-    sourceFilePath?: string
-    testFramework?: string
-    executionCommands?: string[]
-    testCoverage?: number
-    stopIteration?: string
-    errorMessage?: string
-    codeReferences?: References
-    numberOfTestMethods?: number
-}

packages/core/src/codewhisperer/service/testGenHandler.ts

@@ -40,6 +40,7 @@ import fs from '../../shared/fs/fs'
 export function throwIfCancelled() {
     // TODO: fileName will be '' if user gives propt without opening
     if (testGenState.isCancelling()) {
+        testGenState.setToNotStarted()
         throw new TestGenStoppedError()
     }
 }
@@ -150,11 +151,14 @@ export async function pollTestJobStatus(
         logger.debug('pollTestJobStatus testGenerationJob %O', resp.testGenerationJob)
         ChatSessionManager.Instance.getSession().testGenerationJob = resp.testGenerationJob
         const progressRate = resp.testGenerationJob?.progressRate ?? 0
-        testGenState.getChatControllers()?.sendUpdatePromptProgress.fire({
-            tabID: ChatSessionManager.Instance.getSession().tabID,
-            status: 'InProgress',
-            progressRate,
-        })
+        if (initialExecution) {
+            testGenState.getChatControllers()?.sendUpdatePromptProgress.fire({
+                tabID: ChatSessionManager.Instance.getSession().tabID,
+                status: 'InProgress',
+                progressRate,
+            })
+        }
+
         const jobSummary = resp.testGenerationJob?.jobSummary ?? ''
         const jobSummaryNoBackticks = jobSummary.replace(/^`+|`+$/g, '')
         ChatSessionManager.Instance.getSession().jobSummary = jobSummaryNoBackticks
@@ -163,7 +167,8 @@ export async function pollTestJobStatus(
         const targetFileInfo = packageInfo?.targetFileInfoList?.[0]
 
         if (packageInfo) {
-            // TODO: will need some fields from packageInfo such as buildCommand, packagePlan, packageSummary
+            // TODO: will use some fields from packageInfo in the future
+            ChatSessionManager.Instance.getSession().packageInfo = packageInfo
         }
         if (targetFileInfo) {
             if (targetFileInfo.numberOfTestMethods) {
@@ -252,13 +257,11 @@ export async function exportResultsArchive(
                 filePath: testFilePath,
                 projectName,
             })
-
-            // If User accepts the diff
-            testGenState.getChatControllers()?.sendUpdatePromptProgress.fire({
-                tabID: ChatSessionManager.Instance.getSession().tabID,
-                status: 'Completed',
-            })
         }
+        testGenState.getChatControllers()?.sendUpdatePromptProgress.fire({
+            tabID: ChatSessionManager.Instance.getSession().tabID,
+            status: 'Completed',
+        })
     } catch (e) {
         session.numberOfTestsGenerated = 0
         downloadErrorMessage = (e as Error).message