fix(amazonq): downloadExportResultAcrchive api not passed with profileArn and cause 400 error (aws#7022)

## Problem
The CreateUploadUrl and StartTestGenerationJob requests succeeded with
the profile ARN attached. However, the ExportResultArchive requests were
made without this profile ARN, leading to inconsistent profile
resolution between operations.

## Solution


---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: Will-ShaoHua

packages/amazonq/.changes/next-release/Bug Fix-a8339a5e-5367-46a8-82f4-84369b384adb.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fix Q agents will fail for /transform /dev /test features if IdC kms key is configured with 400 error"
+}

packages/core/src/amazonqDoc/session/session.ts

@@ -313,6 +313,7 @@ export class Session {
                     clientId: getClientId(globals.globalState),
                     ideVersion: extensionVersion,
                 },
+                profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
             }
 
             const response = await client.sendTelemetryEvent(params).promise()

packages/core/src/amazonqFeatureDev/client/featureDev.ts

@@ -244,11 +244,13 @@ export class FeatureDevClient implements FeatureClient {
     }
 
     public async exportResultArchive(conversationId: string) {
+        const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
         try {
             const streamingClient = await createCodeWhispererChatStreamingClient()
             const params = {
                 exportId: conversationId,
                 exportIntent: 'TASK_ASSIST',
+                profileArn: profile?.arn,
             } satisfies ExportResultArchiveCommandInput
             getLogger().debug(`Executing exportResultArchive with %O`, params)
             const archiveResponse = await streamingClient.exportResultArchive(params)

packages/core/src/codewhisperer/client/codewhisperer.ts

@@ -252,6 +252,7 @@ export class DefaultCodeWhispererClient {
                 clientId: getClientId(globals.globalState),
                 ideVersion: extensionVersion,
             },
+            profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
         }
         if (!AuthUtil.instance.isValidEnterpriseSsoInUse() && !globals.telemetry.telemetryEnabled) {
             return

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

@@ -72,7 +72,7 @@ export async function startTestGenerationProcess(
         let artifactMap: ArtifactMap = {}
         const uploadStartTime = performance.now()
         try {
-            artifactMap = await getPresignedUrlAndUploadTestGen(zipMetadata)
+            artifactMap = await getPresignedUrlAndUploadTestGen(zipMetadata, profile)
         } finally {
             const outputLogPath = path.join(testGenerationLogsDir, 'output.log')
             if (await fs.existsFile(outputLogPath)) {

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

@@ -269,7 +269,7 @@ export async function preTransformationUploadCode() {
             })
 
             transformByQState.setPayloadFilePath(payloadFilePath)
-            uploadId = await uploadPayload(payloadFilePath)
+            uploadId = await uploadPayload(payloadFilePath, AuthUtil.instance.regionProfileManager.activeRegionProfile)
             telemetry.record({ codeTransformJobId: uploadId }) // uploadId is re-used as jobId
         })
     } catch (err) {
@@ -457,7 +457,7 @@ export async function finishHumanInTheLoop(selectedDependency?: string) {
             }),
         })
 
-        await uploadPayload(uploadResult.tempFilePath, {
+        await uploadPayload(uploadResult.tempFilePath, profile, {
             transformationUploadContext: {
                 jobId,
                 uploadArtifactType: 'Dependencies',

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

@@ -36,6 +36,7 @@ import { randomUUID } from '../../shared/crypto'
 import { sleep } from '../../shared/utilities/timeoutUtils'
 import { tempDirPath } from '../../shared/filesystemUtilities'
 import fs from '../../shared/fs/fs'
+import { AuthUtil } from '../util/authUtil'
 
 // TODO: Get TestFileName and Framework and to error message
 export function throwIfCancelled() {
@@ -45,7 +46,7 @@ export function throwIfCancelled() {
     }
 }
 
-export async function getPresignedUrlAndUploadTestGen(zipMetadata: ZipMetadata) {
+export async function getPresignedUrlAndUploadTestGen(zipMetadata: ZipMetadata, profile: RegionProfile | undefined) {
     const logger = getLogger()
     if (zipMetadata.zipFilePath === '') {
         getLogger().error('Failed to create valid source zip')
@@ -55,6 +56,7 @@ export async function getPresignedUrlAndUploadTestGen(zipMetadata: ZipMetadata)
         contentMd5: getMd5(zipMetadata.zipFilePath),
         artifactType: 'SourceCode',
         uploadIntent: CodeWhispererConstants.testGenUploadIntent,
+        profileArn: profile?.arn,
     }
     logger.verbose(`Prepare for uploading src context...`)
     const srcResp = await codeWhisperer.codeWhispererClient.createUploadUrl(srcReq).catch((err) => {
@@ -310,7 +312,8 @@ export async function downloadResultArchive(
                     },
                 },
             },
-            pathToArchive
+            pathToArchive,
+            AuthUtil.instance.regionProfileManager.activeRegionProfile
         )
     } catch (e: any) {
         downloadErrorMessage = (e as Error).message

packages/core/src/codewhisperer/service/transformByQ/transformApiHandler.ts

@@ -51,6 +51,7 @@ import { encodeHTML } from '../../../shared/utilities/textUtilities'
 import { convertToTimeString } from '../../../shared/datetime'
 import { getAuthType } from '../../../auth/utils'
 import { UserWrittenCodeTracker } from '../../tracker/userWrittenCodeTracker'
+import { AuthUtil } from '../../util/authUtil'
 
 export function getSha256(buffer: Buffer) {
     const hasher = crypto.createHash('sha256')
@@ -193,7 +194,11 @@ export async function stopJob(jobId: string) {
     }
 }
 
-export async function uploadPayload(payloadFileName: string, uploadContext?: UploadContext) {
+export async function uploadPayload(
+    payloadFileName: string,
+    profile: RegionProfile | undefined,
+    uploadContext?: UploadContext
+) {
     const buffer = Buffer.from(await fs.readFileBytes(payloadFileName))
     const sha256 = getSha256(buffer)
 
@@ -205,6 +210,7 @@ export async function uploadPayload(payloadFileName: string, uploadContext?: Upl
             contentChecksumType: CodeWhispererConstants.contentChecksumType,
             uploadIntent: CodeWhispererConstants.uploadIntent,
             uploadContext,
+            profileArn: profile?.arn,
         })
     } catch (e: any) {
         const errorMessage = `Creating the upload URL failed due to: ${(e as Error).message}`
@@ -759,7 +765,8 @@ export async function downloadResultArchive(
                 exportId: jobId,
                 exportIntent: ExportIntent.TRANSFORMATION,
             },
-            pathToArchive
+            pathToArchive,
+            AuthUtil.instance.regionProfileManager.activeRegionProfile
         )
     } catch (e: any) {
         getLogger().error(`CodeTransformation: ExportResultArchive error = %O`, e)

packages/core/src/codewhisperer/service/transformByQ/transformationResultsViewProvider.ts

@@ -27,6 +27,7 @@ import { ChatSessionManager } from '../../../amazonqGumby/chat/storages/chatSess
 import { setContext } from '../../../shared/vscode/setContext'
 import * as codeWhisperer from '../../client/codewhisperer'
 import { UserWrittenCodeTracker } from '../../tracker/userWrittenCodeTracker'
+import { AuthUtil } from '../../util/authUtil'
 
 export abstract class ProposedChangeNode {
     abstract readonly resourcePath: string
@@ -403,7 +404,8 @@ export class ProposedTransformationExplorer {
                             exportId: transformByQState.getJobId(),
                             exportIntent: ExportIntent.TRANSFORMATION,
                         },
-                        pathToArchive
+                        pathToArchive,
+                        AuthUtil.instance.regionProfileManager.activeRegionProfile
                     )
 
                     getLogger().info('CodeTransformation: downloaded results successfully')

packages/core/src/shared/utilities/download.ts

@@ -6,6 +6,7 @@
 import { CodeWhispererStreaming, ExportResultArchiveCommandInput } from '@amzn/codewhisperer-streaming'
 import { ToolkitError } from '../errors'
 import fs from '../fs/fs'
+import { RegionProfile } from '../../codewhisperer/models/model'
 
 /**
  * This class represents the structure of the archive returned by the ExportResultArchive endpoint
@@ -21,9 +22,10 @@ export class ExportResultArchiveStructure {
 export async function downloadExportResultArchive(
     cwStreamingClient: CodeWhispererStreaming,
     exportResultArchiveArgs: ExportResultArchiveCommandInput,
-    toPath: string
+    toPath: string,
+    profile: RegionProfile | undefined
 ) {
-    const result = await cwStreamingClient.exportResultArchive(exportResultArchiveArgs)
+    const result = await cwStreamingClient.exportResultArchive({ ...exportResultArchiveArgs, profileArn: profile?.arn })
 
     const buffer = []