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

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 = []
 

packages/core/src/testInteg/perf/downloadExportResultArchive.test.ts

@@ -11,12 +11,14 @@ import { fs, getRandomString } from '../../shared'
 import { createTestWorkspace } from '../../test/testUtil'
 import { getEqualOSTestOptions, performanceTest } from '../../shared/performance/performance'
 import { downloadExportResultArchive } from '../../shared/utilities/download'
+import { RegionProfile } from '../../codewhisperer'
 
 interface SetupResult {
     workspace: WorkspaceFolder
     exportCommandInput: ExportResultArchiveCommandInput
     writeFileStub: sinon.SinonStub
     cwStreaming: any
+    profile: RegionProfile
 }
 
 interface FakeCommandOutput {
@@ -42,7 +44,8 @@ async function setup(pieces: number, pieceSize: number): Promise<SetupResult> {
     const cwStreaming = { exportResultArchive: () => generateCommandOutput(pieces, pieceSize) }
 
     const writeFileStub = sinon.stub(fs, 'writeFile')
-    return { workspace, exportCommandInput, writeFileStub, cwStreaming }
+    const profile: RegionProfile = { name: 'foo', region: 'us-east-1', arn: 'foo-arn', description: '' }
+    return { workspace, exportCommandInput, writeFileStub, cwStreaming, profile }
 }
 
 function perfTest(pieces: number, pieceSize: number, label: string) {
@@ -56,11 +59,18 @@ function perfTest(pieces: number, pieceSize: number, label: string) {
         function () {
             return {
                 setup: async () => await setup(pieces, pieceSize),
-                execute: async ({ workspace, cwStreaming, exportCommandInput, writeFileStub }: SetupResult) => {
+                execute: async ({
+                    workspace,
+                    cwStreaming,
+                    exportCommandInput,
+                    writeFileStub,
+                    profile,
+                }: SetupResult) => {
                     await downloadExportResultArchive(
                         cwStreaming,
                         exportCommandInput,
-                        path.join(workspace.uri.fsPath, 'result')
+                        path.join(workspace.uri.fsPath, 'result'),
+                        profile
                     )
                 },
                 verify: async (setup: SetupResult) => {