fix error handling

Author: zixlin7

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

@@ -217,7 +217,7 @@ class ChatController private constructor(
 
                     editor.document.insertString(offset, message.code)
 
-                    ReferenceLogController.addReferenceLog(message.code, message.codeReference, editor, context.project)
+                    ReferenceLogController.addReferenceLog(message.code, message.codeReference, editor, context.project, null)
 
                     CodeWhispererUserModificationTracker.getInstance(context.project).enqueue(
                         InsertedCodeModificationEntry(

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

@@ -9,11 +9,12 @@ 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.model.CaretPosition
 import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
 import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
 
 object ReferenceLogController {
-    fun addReferenceLog(originalCode: String, codeReferences: List<CodeReference>?, editor: Editor, project: Project) {
+    fun addReferenceLog(originalCode: String, codeReferences: List<CodeReference>?, editor: Editor, project: Project, inlineChatStartPosition: CaretPosition?) {
         codeReferences?.let { references ->
             val cwReferences = references.map { reference ->
                 Reference.builder()
@@ -36,7 +37,7 @@ object ReferenceLogController {
                 originalCode,
                 cwReferences,
                 editor,
-                CodeWhispererEditorUtil.getCaretPosition(editor),
+                inlineChatStartPosition ?: CodeWhispererEditorUtil.getCaretPosition(editor),
                 null,
             )
         }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatController.kt

@@ -57,8 +57,10 @@ import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AMAZON_Q_WINDOW_ID
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
+import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.ChatRequestData
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.TriggerType
+import software.aws.toolkits.jetbrains.services.cwc.controller.ReferenceLogController
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.messenger.ChatPromptHandler
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.TelemetryHelper
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.userIntent.UserIntentRecognizer
@@ -68,6 +70,7 @@ import software.aws.toolkits.jetbrains.services.cwc.inline.listeners.InlineChatF
 import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessage
 import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType
 import software.aws.toolkits.jetbrains.services.cwc.storage.ChatSessionStorage
+import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
 import software.aws.toolkits.telemetry.FeatureId
 import java.util.UUID
 import java.util.concurrent.atomic.AtomicBoolean
@@ -285,9 +288,10 @@ class InlineChatController(
         .replace("'", "'")
         .replace("=>", "=>")
 
-    private fun processNewCode(editor: Editor, line: Int, code: String, prevMessage: String) {
+    private fun processNewCode(editor: Editor, line: Int, event: ChatMessage, prevMessage: String) {
         if (isAbandoned.get()) return
         runBlocking {
+            val code = event.message?.let { unescape(it) } ?: return@runBlocking
             logger.debug { "received inline chat recommendation with code: \n $code" }
             var insertLine = line
             var linesToAdd = emptyList<String>()
@@ -315,6 +319,12 @@ class InlineChatController(
             insertionLine.set(insertLine)
             acceptAction = {
                 removeHighlighter(editor)
+                try {
+                    val caretPosition = CaretPosition(offset = getLineStartOffset(editor.document, line), line = line)
+                    ReferenceLogController.addReferenceLog(code, event.codeReference, editor, project, inlineChatStartPosition = caretPosition)
+                } catch (e: Exception) {
+                    logger.warn { "error logging reference for inline chat: ${e.stackTraceToString()}" }
+                }
             }
             rejectAction = {
                 val startOffset = getLineStartOffset(editor.document, line)
@@ -472,6 +482,13 @@ class InlineChatController(
                     val endOffset = getLineEndOffset(editor.document, endLine)
                     replaceString(editor.document, startOffset, endOffset, recommendation)
                     removeHighlighter(editor)
+                    try {
+                        val caretPosition =
+                            CaretPosition(offset = selectionRange.startOffset, line = getLineNumber(editor.document, selectionRange.startOffset))
+                        ReferenceLogController.addReferenceLog(recommendation, event.codeReference, editor, project, inlineChatStartPosition = caretPosition)
+                    } catch (e: Exception) {
+                        logger.warn { "error logging reference for inline chat: ${e.stackTraceToString()}" }
+                    }
                 }
                 rejectAction = {
                     val startOffset = getLineStartOffset(editor.document, startLine)
@@ -624,6 +641,7 @@ class InlineChatController(
 
         val sessionInfo = sessionStorage.getSession(tabId, project)
         val mutex = Mutex()
+        val isReferenceAllowed = CodeWhispererSettings.getInstance().isIncludeCodeWithReference()
 
         var errorMessage = ""
         var prevMessage = ""
@@ -645,11 +663,17 @@ class InlineChatController(
                 .onEach { event: ChatMessage ->
                     if (event.message?.isNotEmpty() == true && prevMessage != event.message) {
                         mutex.withLock {
+                            if (event.codeReference?.isNotEmpty() == true && !isReferenceAllowed) {
+                                canPopupAbort.set(true)
+                                undoChanges()
+                                errorMessage = "Suggestion had code reference; removed per setting."
+                                return@withLock
+                            }
                             try {
                                 if (selectionRange != null) {
                                     processChatDiff(selectedCode, event, editor, selectionRange!!)
                                 } else {
-                                    processNewCode(editor, selectedLineStart, unescape(event.message), prevMessage)
+                                    processNewCode(editor, selectedLineStart, event, prevMessage)
                                 }
                             } catch (e: Exception) {
                                 canPopupAbort.set(true)

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatPopupFactory.kt

@@ -70,8 +70,10 @@ class InlineChatPopupFactory(
                             errorMessage = submitHandler(prompt, selectedCode, selectedLineStart, editor)
                         }
                         if (errorMessage.isNotEmpty()) {
-                            setLabel(errorMessage)
-                            revalidate()
+                            withContext(EDT) {
+                                setErrorMessage(errorMessage)
+                                revalidate()
+                            }
                         } else {
                             val acceptAction = {
                                 acceptHandler.invoke()
@@ -81,6 +83,7 @@ class InlineChatPopupFactory(
                             }
                             withContext(EDT) {
                                 addCodeActionsPanel(acceptAction, rejectAction)
+                                revalidate()
                             }
                         }
                     }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatPopupPanel.kt

@@ -88,6 +88,10 @@ class InlineChatPopupPanel(private val parentDisposable: Disposable) : JPanel()
         add(buttonsPanel, BorderLayout.EAST)
     }
 
+    private val emptyTextField = createTextField().apply {
+        isEnabled = false
+    }
+
     override fun getPreferredSize(): Dimension = Dimension(popupWidth, popupHeight)
 
     private fun createTextField(): JTextField = JTextField().apply {
@@ -141,9 +145,7 @@ class InlineChatPopupPanel(private val parentDisposable: Disposable) : JPanel()
     fun addCodeActionsPanel(acceptAction: () -> Unit, rejectAction: () -> Unit) {
         textLabel.text = message("amazonqInlineChat.popup.editCode")
         // this is a workaround somehow the textField will interfere with the enter handler
-        val emptyTextField = createTextField()
         emptyTextField.text = textField.text
-        emptyTextField.isEnabled = false
         inputPanel.remove(textField)
         inputPanel.add(emptyTextField, BorderLayout.CENTER)
 
@@ -162,6 +164,22 @@ class InlineChatPopupPanel(private val parentDisposable: Disposable) : JPanel()
         revalidate()
     }
 
+    fun setErrorMessage(message: String) {
+        setLabel(message)
+        inputPanel.remove(emptyTextField)
+        textField.text = ""
+        textField.isEnabled = true
+        inputPanel.add(textField, BorderLayout.CENTER)
+
+        buttonsPanel.remove(acceptButton)
+        buttonsPanel.remove(rejectButton)
+        buttonsPanel.add(submitButton, BorderLayout.WEST)
+        buttonsPanel.add(cancelButton, BorderLayout.EAST)
+        submitButton.isEnabled = false
+        cancelButton.isEnabled = true
+        revalidate()
+    }
+
     fun setLabel(text: String) {
         textLabel.text = text
         textField.isEnabled = false