Fixed feature dev code generated code parsing to handle references (aws#4414)

* Fixed feature dev code generated code parsing to handle references

* added changelog

* addressed comments

Author: kumsmrit

.changes/next-release/bugfix-85cf787d-7de8-4d3c-a0a2-5a54632f4ff3.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Amazon Q Feature Development: Handle generated code parsing for rendering references correctly"
+}

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

@@ -54,14 +54,14 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendM
 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.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.SessionStatePhase
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
-import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
 import software.aws.toolkits.jetbrains.ui.feedback.FeatureDevFeedbackDialog
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.AmazonqTelemetry
@@ -279,7 +279,7 @@ class FeatureDevController(
 
             var filePaths: List<NewFileZipInfo> = emptyList()
             var deletedFiles: List<DeletedFileInfo> = emptyList()
-            var references: List<CodeReference> = emptyList()
+            var references: List<CodeReferenceGenerated> = emptyList()
 
             when (val state = session.sessionState) {
                 is PrepareCodeGenerationState -> {

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

@@ -14,12 +14,12 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendC
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendCodeResult
 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.session.CodeReferenceGenerated
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo
 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.util.getFollowUpOptions
-import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.resources.message
 
@@ -45,7 +45,7 @@ suspend fun FeatureDevController.onCodeGeneration(session: Session, message: Str
 
         var filePaths: List<NewFileZipInfo> = emptyList()
         var deletedFiles: List<DeletedFileInfo> = emptyList()
-        var references: List<CodeReference> = emptyList()
+        var references: List<CodeReferenceGenerated> = emptyList()
         var uploadId = ""
 
         when (state) {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessage.kt

@@ -182,7 +182,7 @@ data class CodeResultMessage(
     val conversationId: String,
     val filePaths: List<NewFileZipInfo>,
     val deletedFiles: List<DeletedFileInfo>,
-    val references: List<ReducedCodeReference>
+    val references: List<CodeReference>
 ) : UiMessage(
     tabId = tabId,
     type = "codeResultMessage"

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessagePublisherExtensions.kt

@@ -5,11 +5,12 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages
 
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
+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.NewFileZipInfo
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase
-import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.licenseText
 import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
+import software.aws.toolkits.jetbrains.services.cwc.messages.RecommendationContentSpan
 import software.aws.toolkits.resources.message
 import java.util.UUID
 
@@ -191,11 +192,19 @@ suspend fun MessagePublisher.sendCodeResult(
     uploadId: String,
     filePaths: List<NewFileZipInfo>,
     deletedFiles: List<DeletedFileInfo>,
-    references: List<CodeReference>
+    references: List<CodeReferenceGenerated>
 ) {
-    // It is being done this mapping as featureDev currently doesn't support fully references.
-    val refs = references.filter { it.licenseName != null && it.url != null && it.repository != null }.map {
-        ReducedCodeReference(information = it.licenseText())
+    val refs = references.map { ref ->
+        CodeReference(
+            licenseName = ref.licenseName,
+            repository = ref.repository,
+            url = ref.url,
+            recommendationContentSpan = RecommendationContentSpan(
+                ref.recommendationContentSpan?.start ?: 0,
+                ref.recommendationContentSpan?.end ?: 0,
+            ),
+            information = "Reference code under **${ref.licenseName}** license from repository [${ref.repository}](${ref.url})"
+        )
     }
 
     this.publish(

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

@@ -8,7 +8,6 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendA
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.createUploadUrl
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.deleteUploadArtifact
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3
-import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.AmazonqTelemetry
 import software.aws.toolkits.telemetry.AmazonqUploadIntent
@@ -19,7 +18,7 @@ class PrepareCodeGenerationState(
     private var config: SessionStateConfig,
     val filePaths: List<NewFileZipInfo>,
     val deletedFiles: List<DeletedFileInfo>,
-    val references: List<CodeReference>,
+    val references: List<CodeReferenceGenerated>,
     var uploadId: String,
     private val currentIteration: Int,
     private var messenger: MessagePublisher

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

@@ -16,7 +16,6 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.createCon
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.resolveAndCreateOrUpdateFile
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.resolveAndDeleteFile
 import software.aws.toolkits.jetbrains.services.cwc.controller.ReferenceLogController
-import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
 import software.aws.toolkits.telemetry.AmazonqTelemetry
 
 class Session(val tabID: String, val project: Project) {
@@ -93,7 +92,7 @@ class Session(val tabID: String, val project: Project) {
     /**
      * Triggered by the Insert code follow-up button to apply code changes.
      */
-    fun insertChanges(filePaths: List<NewFileZipInfo>, deletedFiles: List<DeletedFileInfo>, references: List<CodeReference>) {
+    fun insertChanges(filePaths: List<NewFileZipInfo>, deletedFiles: List<DeletedFileInfo>, references: List<CodeReferenceGenerated>) {
         val projectRootPath = context.projectRoot.toNioPath()
 
         filePaths.forEach { resolveAndCreateOrUpdateFile(projectRootPath, it.zipFilePath, it.fileContent) }

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

@@ -6,7 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 import com.fasterxml.jackson.annotation.JsonValue
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient
-import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
+import software.aws.toolkits.jetbrains.services.cwc.messages.RecommendationContentSpan
 
 data class SessionStateAction(
     val task: String,
@@ -51,14 +51,21 @@ data class DeletedFileInfo(
 data class CodeGenerationResult(
     var newFiles: List<NewFileZipInfo>,
     var deletedFiles: List<DeletedFileInfo>,
-    var references: List<CodeReference>,
+    var references: List<CodeReferenceGenerated>,
+)
+
+data class CodeReferenceGenerated(
+    val licenseName: String? = null,
+    val repository: String? = null,
+    val url: String? = null,
+    val recommendationContentSpan: RecommendationContentSpan? = null,
 )
 
 @Suppress("ConstructorParameterNaming") // Unfortunately, this is exactly how the string json is received and is needed for parsing.
 data class CodeGenerationStreamResult(
     var new_file_contents: Map<String, String>,
     var deleted_files: List<String>,
-    var references: List<CodeReference>,
+    var references: List<CodeReferenceGenerated>,
 )
 
 @Suppress("ConstructorParameterNaming") // Unfortunately, this is exactly how the string json is received and is needed for parsing.

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevClientUtil.kt

@@ -14,6 +14,7 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.CodeWhispere
 import software.amazon.awssdk.services.codewhispererstreaming.model.ServiceQuotaExceededException
 import software.amazon.awssdk.services.codewhispererstreaming.model.ThrottlingException
 import software.aws.toolkits.core.utils.debug
+import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeIterationLimitError
@@ -207,6 +208,7 @@ suspend fun exportTaskAssistArchiveResult(proxyClient: FeatureDevClient, convers
         val result = exportResponse.reduce { acc, next -> acc + next } // To map the result it is needed to combine the  full byte array
         parsedResult = jacksonObjectMapper().readValue(result)
     } catch (e: Exception) {
+        logger.error(e) { "Failed to parse downloaded code results" }
         exportParseError()
     }
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ReferenceLogController.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import software.amazon.awssdk.services.codewhispererruntime.model.Reference
 import software.amazon.awssdk.services.codewhispererruntime.model.Span
+import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeReferenceGenerated
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
 import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
 import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
@@ -41,7 +42,7 @@ object ReferenceLogController {
         }
     }
 
-    fun addReferenceLog(codeReferences: List<CodeReference>?, project: Project) {
+    fun addReferenceLog(codeReferences: List<CodeReferenceGenerated>?, project: Project) {
         val manager = CodeWhispererCodeReferenceManager.getInstance(project)
 
         codeReferences?.forEach { reference ->