From 89b59302fcc6f0e1c0affb6671ecc28ad13df2e1 Mon Sep 17 00:00:00 2001
From: Laxman Reddy <141967714+laileni-aws@users.noreply.github.com>
Date: Wed, 11 Dec 2024 11:05:49 -0800
Subject: [PATCH 1/2] Removing session cleanUp before emitting the telemetry
 event

---
 packages/core/src/amazonqTest/chat/controller/controller.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/packages/core/src/amazonqTest/chat/controller/controller.ts b/packages/core/src/amazonqTest/chat/controller/controller.ts
index 88490315e45..66480273030 100644
--- a/packages/core/src/amazonqTest/chat/controller/controller.ts
+++ b/packages/core/src/amazonqTest/chat/controller/controller.ts
@@ -317,7 +317,6 @@ export class TestController {
             case ButtonActions.STOP_TEST_GEN:
                 testGenState.setToCancelling()
                 telemetry.ui_click.emit({ elementId: 'unitTestGeneration_cancelTestGenerationProgress' })
-                await this.sessionCleanUp()
                 break
             case ButtonActions.STOP_BUILD:
                 cancelBuild()

From cf81315a86e8bb5369aeb88487f612f1cf200149 Mon Sep 17 00:00:00 2001
From: Laxman Reddy <141967714+laileni-aws@users.noreply.github.com>
Date: Mon, 17 Feb 2025 09:22:30 -0800
Subject: [PATCH 2/2] Adding requestId for uploadArtifactToS3 for testGen

---
 .../core/src/codewhisperer/service/securityScanHandler.ts   | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/codewhisperer/service/securityScanHandler.ts b/packages/core/src/codewhisperer/service/securityScanHandler.ts
index cf344610e27..1a22ff04701 100644
--- a/packages/core/src/codewhisperer/service/securityScanHandler.ts
+++ b/packages/core/src/codewhisperer/service/securityScanHandler.ts
@@ -39,12 +39,13 @@ import {
     SecurityScanTimedOutError,
     UploadArtifactToS3Error,
 } from '../models/errors'
-import { getTelemetryReasonDesc } from '../../shared/errors'
+import { getTelemetryReasonDesc, isAwsError } from '../../shared/errors'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { detectCommentAboveLine } from '../../shared/utilities/commentUtils'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
 import { FeatureUseCase } from '../models/constants'
 import { UploadTestArtifactToS3Error } from '../../amazonqTest/error'
+import { ChatSessionManager } from '../../amazonqTest/chat/storages/chatSession'
 
 export async function listScanResults(
     client: DefaultCodeWhispererClient,
@@ -386,6 +387,9 @@ export async function uploadArtifactToS3(
         } else {
             errorMessage = errorDesc ?? defaultMessage
         }
+        if (isAwsError(error) && featureUseCase === FeatureUseCase.TEST_GENERATION) {
+            ChatSessionManager.Instance.getSession().startTestGenerationRequestId = error.requestId
+        }
         throw isCodeScan ? new UploadArtifactToS3Error(errorMessage) : new UploadTestArtifactToS3Error(errorMessage)
     }
 }