diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt index 1dfd5d556be..03426f021b5 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt @@ -46,6 +46,7 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererCon import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage +import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants @@ -92,6 +93,7 @@ interface CodeWhispererClientAdaptor : Disposable { fun listAvailableCustomizations(): List fun sendUserTriggerDecisionTelemetry( + sessionContext: SessionContext, requestContext: RequestContext, responseContext: ResponseContext, completionType: CodewhispererCompletionType, @@ -293,6 +295,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } override fun sendUserTriggerDecisionTelemetry( + sessionContext: SessionContext, requestContext: RequestContext, responseContext: ResponseContext, completionType: CodewhispererCompletionType, @@ -303,24 +306,24 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW ): SendTelemetryEventResponse { val fileContext = requestContext.fileContextInfo val programmingLanguage = fileContext.programmingLanguage - var e2eLatency = requestContext.latencyContext.getCodeWhispererEndToEndLatency() + var e2eLatency = sessionContext.latencyContext.getCodeWhispererEndToEndLatency() // When we send a userTriggerDecision of Empty or Discard, we set the time users see the first // suggestion to be now. if (e2eLatency < 0) { e2eLatency = TimeUnit.NANOSECONDS.toMillis( - System.nanoTime() - requestContext.latencyContext.codewhispererEndToEndStart + System.nanoTime() - sessionContext.latencyContext.codewhispererEndToEndStart ).toDouble() } return bearerClient().sendTelemetryEvent { requestBuilder -> requestBuilder.telemetryEvent { telemetryEventBuilder -> telemetryEventBuilder.userTriggerDecisionEvent { - it.requestId(requestContext.latencyContext.firstRequestId) + it.requestId(sessionContext.latencyContext.firstRequestId) it.completionType(completionType.toCodeWhispererSdkType()) it.programmingLanguage { builder -> builder.languageName(programmingLanguage.toCodeWhispererRuntimeLanguage().languageId) } it.sessionId(responseContext.sessionId) it.recommendationLatencyMilliseconds(e2eLatency) - it.triggerToResponseLatencyMilliseconds(requestContext.latencyContext.paginationFirstCompletionTime) + it.triggerToResponseLatencyMilliseconds(sessionContext.latencyContext.paginationFirstCompletionTime) it.suggestionState(suggestionState.toCodeWhispererSdkType()) it.timestamp(Instant.now()) it.suggestionReferenceCount(suggestionReferenceCount) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt index 22c8c2176bf..62f1b2f47a7 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt @@ -11,9 +11,9 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.util.TextRange import com.intellij.psi.PsiDocumentManager import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition -import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.PAIRED_BRACKETS @@ -23,17 +23,21 @@ import java.util.Stack @Service class CodeWhispererEditorManager { - fun updateEditorWithRecommendation(states: InvocationContext, sessionContext: SessionContext) { - val (requestContext, responseContext, recommendationContext) = states - val (project, editor) = requestContext + fun updateEditorWithRecommendation(sessionContext: SessionContext) { + val previews = CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo() + val selectedIndex = sessionContext.selectedIndex + val preview = previews[selectedIndex] + val states = CodeWhispererService.getInstance().getAllPaginationSessions()[preview.jobId] ?: return + val (requestContext, responseContext) = states + val (project, editor) = sessionContext val document = editor.document val primaryCaret = editor.caretModel.primaryCaret - val selectedIndex = sessionContext.selectedIndex - val typeahead = sessionContext.typeahead - val detail = recommendationContext.details[selectedIndex] + val typeahead = preview.typeahead + val detail = preview.detail + val userInput = preview.userInput val reformatted = CodeWhispererPopupManager.getInstance().getReformattedRecommendation( detail, - recommendationContext.userInputSinceInvocation + userInput ) val remainingRecommendation = reformatted.substring(typeahead.length) val originalOffset = primaryCaret.offset - typeahead.length @@ -43,6 +47,8 @@ class CodeWhispererEditorManager { val insertEndOffset = sessionContext.insertEndOffset val endOffsetToReplace = if (insertEndOffset != -1) insertEndOffset else primaryCaret.offset + preview.detail.isAccepted = true + WriteCommandAction.runWriteCommandAction(project) { document.replaceString(originalOffset, endOffsetToReplace, reformatted) PsiDocumentManager.getInstance(project).commitDocument(document) @@ -67,7 +73,7 @@ class CodeWhispererEditorManager { ApplicationManager.getApplication().messageBus.syncPublisher( CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED, - ).afterAccept(states, sessionContext, rangeMarker) + ).afterAccept(states, previews, sessionContext, rangeMarker) } } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt index e699e97e856..6fec203cffb 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt @@ -95,7 +95,7 @@ object CodeWhispererEditorUtil { } fun shouldSkipInvokingBasedOnRightContext(editor: Editor): Boolean { - val caretContext = runReadAction { CodeWhispererEditorUtil.extractCaretContext(editor) } + val caretContext = runReadAction { extractCaretContext(editor) } val rightContextLines = caretContext.rightFileContext.split(Regex("\r?\n")) val rightContextCurrentLine = if (rightContextLines.isEmpty()) "" else rightContextLines[0] diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdder.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdder.kt index ce56ea8fc60..6c3bd789c2e 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdder.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdder.kt @@ -14,15 +14,15 @@ import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext abstract class CodeWhispererImportAdder { abstract val supportedLanguages: List abstract val dummyFileName: String - fun insertImportStatements(states: InvocationContext, sessionContext: SessionContext) { - val imports = states.recommendationContext.details[sessionContext.selectedIndex] - .recommendation.mostRelevantMissingImports() + fun insertImportStatements(states: InvocationContext, previews: List, sessionContext: SessionContext) { + val imports = previews[sessionContext.selectedIndex].detail.recommendation.mostRelevantMissingImports() LOG.info { "Adding ${imports.size} imports for completions, sessionId: ${states.responseContext.sessionId}" } imports.forEach { insertImportStatement(states, it) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdderListener.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdderListener.kt index 0a2f8110b5d..a583bd54cde 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdderListener.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdderListener.kt @@ -7,13 +7,14 @@ import com.intellij.openapi.editor.RangeMarker import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings object CodeWhispererImportAdderListener : CodeWhispererUserActionListener { internal val LOG = getLogger() - override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) { + override fun afterAccept(states: InvocationContext, previews: List, sessionContext: SessionContext, rangeMarker: RangeMarker) { if (!CodeWhispererSettings.getInstance().isImportAdderEnabled()) { LOG.debug { "Import adder not enabled in user settings" } return @@ -28,6 +29,6 @@ object CodeWhispererImportAdderListener : CodeWhispererUserActionListener { LOG.debug { "No import adder found for $language" } return } - importAdder.insertImportStatements(states, sessionContext) + importAdder.insertImportStatements(states, previews, sessionContext) } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/inlay/CodeWhispererInlayManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/inlay/CodeWhispererInlayManager.kt index e72269ce2f7..60daf84c948 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/inlay/CodeWhispererInlayManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/inlay/CodeWhispererInlayManager.kt @@ -6,28 +6,26 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.inlay import com.intellij.idea.AppMode import com.intellij.openapi.components.Service import com.intellij.openapi.components.service -import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.EditorCustomElementRenderer import com.intellij.openapi.editor.Inlay -import com.intellij.openapi.ui.popup.JBPopup import com.intellij.openapi.util.Disposer -import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationChunk +import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext @Service class CodeWhispererInlayManager { private val existingInlays = mutableListOf>() - fun updateInlays(states: InvocationContext, chunks: List) { - val editor = states.requestContext.editor + fun updateInlays(sessionContext: SessionContext, chunks: List) { clearInlays() chunks.forEach { chunk -> - createCodeWhispererInlays(editor, chunk.inlayOffset, chunk.text, states.popup) + createCodeWhispererInlays(sessionContext, chunk.inlayOffset, chunk.text) } } - private fun createCodeWhispererInlays(editor: Editor, startOffset: Int, inlayText: String, popup: JBPopup) { + private fun createCodeWhispererInlays(sessionContext: SessionContext, startOffset: Int, inlayText: String) { if (inlayText.isEmpty()) return + val editor = sessionContext.editor val firstNewlineIndex = inlayText.indexOf("\n") val firstLine: String val otherLines: String @@ -49,7 +47,7 @@ class CodeWhispererInlayManager { val inlineInlay = editor.inlayModel.addInlineElement(startOffset, true, firstLineRenderer) inlineInlay?.let { existingInlays.add(it) - Disposer.register(popup, it) + Disposer.register(sessionContext, it) } } @@ -73,7 +71,7 @@ class CodeWhispererInlayManager { ) blockInlay?.let { existingInlays.add(it) - Disposer.register(popup, it) + Disposer.register(sessionContext, it) } } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt index cf465134162..5c3f53edcdb 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt @@ -4,18 +4,24 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.model import com.intellij.openapi.Disposable +import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.VisualPosition import com.intellij.openapi.editor.markup.RangeHighlighter +import com.intellij.openapi.project.Project import com.intellij.openapi.ui.popup.JBPopup +import com.intellij.openapi.util.Disposer import com.intellij.openapi.vfs.VirtualFile +import com.intellij.util.concurrency.annotations.RequiresEdt import software.amazon.awssdk.services.codewhispererruntime.model.Completion import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext +import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants import software.aws.toolkits.jetbrains.services.codewhisperer.util.CrossFileStrategy import software.aws.toolkits.jetbrains.services.codewhisperer.util.SupplementalContextStrategy @@ -23,6 +29,8 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.UtgStrategy import software.aws.toolkits.telemetry.CodewhispererCompletionType import software.aws.toolkits.telemetry.CodewhispererTriggerType import software.aws.toolkits.telemetry.Result +import java.time.Duration +import java.time.Instant import java.util.concurrent.TimeUnit data class Chunk( @@ -81,10 +89,19 @@ data class SupplementalContextInfo( } data class RecommendationContext( - val details: List, + val details: MutableList, val userInputOriginal: String, val userInputSinceInvocation: String, val position: VisualPosition, + val jobId: Int, + var typeahead: String = "", +) + +data class PreviewContext( + val jobId: Int, + val detail: DetailContext, + val userInput: String, + val typeahead: String, ) data class DetailContext( @@ -95,17 +112,47 @@ data class DetailContext( val isTruncatedOnRight: Boolean, val rightOverlap: String = "", val completionType: CodewhispererCompletionType, + var hasSeen: Boolean = false, + var isAccepted: Boolean = false ) data class SessionContext( - val typeahead: String = "", - val typeaheadOriginal: String = "", - val selectedIndex: Int = 0, + val project: Project, + val editor: Editor, + var popup: JBPopup? = null, + var selectedIndex: Int = -1, val seen: MutableSet = mutableSetOf(), - val isFirstTimeShowingPopup: Boolean = true, + var isFirstTimeShowingPopup: Boolean = true, var toBeRemovedHighlighter: RangeHighlighter? = null, var insertEndOffset: Int = -1, -) + var popupOffset: Int = -1, + val latencyContext: LatencyContext, + var hasAccepted: Boolean = false +) : Disposable { + private var isDisposed = false + + @RequiresEdt + override fun dispose() { + CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll( + this, + hasAccepted, + CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) } + ) + CodeWhispererInvocationStatus.getInstance().setDisplaySessionActive(false) + + if (hasAccepted) { + popup?.closeOk(null) + } else { + popup?.cancel() + } + popup?.let { Disposer.dispose(it) } + popup = null + CodeWhispererInvocationStatus.getInstance().finishInvocation() + isDisposed = true + } + + fun isDisposed() = isDisposed +} data class RecommendationChunk( val text: String, @@ -124,16 +171,21 @@ data class InvocationContext( val requestContext: RequestContext, val responseContext: ResponseContext, val recommendationContext: RecommendationContext, - val popup: JBPopup, ) : Disposable { - override fun dispose() {} + private var isDisposed = false + + @RequiresEdt + override fun dispose() { + isDisposed = true + } + + fun isDisposed() = isDisposed } data class WorkerContext( val requestContext: RequestContext, val responseContext: ResponseContext, val response: GenerateCompletionsResponse, - val popup: JBPopup, ) data class CodeScanTelemetryEvent( diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt index c485c52acbe..8d4e288b5f7 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt @@ -5,30 +5,11 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.popup import com.intellij.openapi.ui.popup.JBPopupListener import com.intellij.openapi.ui.popup.LightweightWindowEvent -import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext -import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus -import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService -import java.time.Duration -import java.time.Instant +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService -class CodeWhispererPopupListener(private val states: InvocationContext) : JBPopupListener { - override fun beforeShown(event: LightweightWindowEvent) { - super.beforeShown(event) - CodeWhispererInvocationStatus.getInstance().setPopupStartTimestamp() - } +class CodeWhispererPopupListener : JBPopupListener { override fun onClosed(event: LightweightWindowEvent) { super.onClosed(event) - val (requestContext, responseContext, recommendationContext) = states - - CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll( - requestContext, - responseContext, - recommendationContext, - CodeWhispererPopupManager.getInstance().sessionContext, - event.isOk, - CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) } - ) - - CodeWhispererInvocationStatus.getInstance().setPopupActive(false) + CodeWhispererService.getInstance().disposeDisplaySession(event.isOk) } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt index c548dff1d4d..1bdd36ee42b 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt @@ -29,16 +29,15 @@ import com.intellij.openapi.editor.event.DocumentEvent import com.intellij.openapi.editor.event.DocumentListener import com.intellij.openapi.editor.event.SelectionEvent import com.intellij.openapi.editor.event.SelectionListener -import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.openapi.project.Project import com.intellij.openapi.ui.popup.JBPopup import com.intellij.openapi.ui.popup.JBPopupFactory import com.intellij.openapi.util.Disposer -import com.intellij.openapi.wm.WindowManager import com.intellij.ui.ComponentUtil import com.intellij.ui.awt.RelativePoint import com.intellij.ui.popup.AbstractPopup import com.intellij.ui.popup.PopupFactoryImpl +import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.messages.Topic import com.intellij.util.ui.UIUtil import software.amazon.awssdk.services.codewhispererruntime.model.Import @@ -51,6 +50,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhisper import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.inlineLabelConstraints import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.popup.handlers.CodeWhispererEditorActionHandler import software.aws.toolkits.jetbrains.services.codewhisperer.popup.handlers.CodeWhispererPopupBackspaceHandler @@ -65,7 +65,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.Co import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererPrevButtonActionListener import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererScrollListener import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus -import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.POPUP_INFO_TEXT_SIZE @@ -85,10 +85,6 @@ class CodeWhispererPopupManager { var shouldListenerCancelPopup: Boolean = true private set - var sessionContext = SessionContext() - private set - - private var myPopup: JBPopup? = null init { // Listen for global scheme changes @@ -115,113 +111,107 @@ class CodeWhispererPopupManager { ) } - fun changeStates( - states: InvocationContext, - indexChange: Int, - typeaheadChange: String, - typeaheadAdded: Boolean, - recommendationAdded: Boolean = false, - ) { - val (_, _, recommendationContext, popup) = states - val (details) = recommendationContext - if (recommendationAdded) { - LOG.debug { - "Add recommendations to the existing CodeWhisperer session, current number of recommendations: ${details.size}" - } - ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_POPUP_STATE_CHANGED) - .recommendationAdded(states, sessionContext) - return - } - val typeaheadOriginal = - if (typeaheadAdded) { - sessionContext.typeaheadOriginal + typeaheadChange - } else { - if (typeaheadChange.length > sessionContext.typeaheadOriginal.length) { - cancelPopup(popup) - return - } - sessionContext.typeaheadOriginal.substring( - 0, - sessionContext.typeaheadOriginal.length - typeaheadChange.length - ) - } - val isReverse = indexChange < 0 - val userInput = states.recommendationContext.userInputSinceInvocation - val validCount = getValidCount(details, userInput, typeaheadOriginal) - val validSelectedIndex = getValidSelectedIndex(details, userInput, sessionContext.selectedIndex, typeaheadOriginal) + @RequiresEdt + fun changeStatesForNavigation(sessionContext: SessionContext, indexChange: Int) { + val validCount = getValidCount() + val validSelectedIndex = getValidSelectedIndex(sessionContext.selectedIndex) if ((validSelectedIndex == validCount - 1 && indexChange == 1) || (validSelectedIndex == 0 && indexChange == -1) ) { return } - val selectedIndex = findNewSelectedIndex( - isReverse, - details, - userInput, - sessionContext.selectedIndex + indexChange, - typeaheadOriginal + val isReverse = indexChange < 0 + val selectedIndex = findNewSelectedIndex(isReverse, sessionContext.selectedIndex + indexChange) + + sessionContext.selectedIndex = selectedIndex + sessionContext.isFirstTimeShowingPopup = false + + ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_POPUP_STATE_CHANGED).stateChanged( + sessionContext ) - if (selectedIndex == -1 || !isValidRecommendation(details[selectedIndex], userInput, typeaheadOriginal)) { - LOG.debug { "None of the recommendation is valid at this point, cancelling the popup" } - cancelPopup(popup) + } + + @RequiresEdt + fun changeStatesForTypeahead( + sessionContext: SessionContext, + typeaheadChange: String, + typeaheadAdded: Boolean + ) { + updateTypeahead(typeaheadChange, typeaheadAdded) + updateSessionSelectedIndex(sessionContext) + sessionContext.isFirstTimeShowingPopup = false + + ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_POPUP_STATE_CHANGED).stateChanged( + sessionContext + ) + } + + @RequiresEdt + fun changeStatesForShowing(sessionContext: SessionContext, states: InvocationContext, recommendationAdded: Boolean = false) { + sessionContext.isFirstTimeShowingPopup = !recommendationAdded + if (recommendationAdded) { + ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_POPUP_STATE_CHANGED) + .recommendationAdded(states, sessionContext) return } - val typeahead = resolveTypeahead(states, selectedIndex, typeaheadOriginal) - val isFirstTimeShowingPopup = indexChange == 0 && typeaheadChange.isEmpty() - sessionContext = SessionContext( - typeahead, - typeaheadOriginal, - selectedIndex, - sessionContext.seen, - isFirstTimeShowingPopup, - sessionContext.toBeRemovedHighlighter - ) + + updateSessionSelectedIndex(sessionContext) + if (sessionContext.popupOffset == -1) { + sessionContext.popupOffset = sessionContext.editor.caretModel.offset + } ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_POPUP_STATE_CHANGED).stateChanged( - states, sessionContext ) } - private fun resolveTypeahead(states: InvocationContext, selectedIndex: Int, typeahead: String): String { - val recommendation = states.recommendationContext.details[selectedIndex].reformatted.content() - val userInput = states.recommendationContext.userInputSinceInvocation - var indexOfFirstNonWhiteSpace = typeahead.indexOfFirst { !it.isWhitespace() } - if (indexOfFirstNonWhiteSpace == -1) { - indexOfFirstNonWhiteSpace = typeahead.length + private fun updateTypeahead(typeaheadChange: String, typeaheadAdded: Boolean) { + val recommendations = CodeWhispererService.getInstance().getAllPaginationSessions().values.filterNotNull() + recommendations.forEach { + val newTypeahead = + if (typeaheadAdded) { + it.recommendationContext.typeahead + typeaheadChange + } else { + if (typeaheadChange.length > it.recommendationContext.typeahead.length) { + LOG.debug { "Typeahead change is longer than the current typeahead, exiting the session" } + CodeWhispererService.getInstance().disposeDisplaySession(false) + return + } + it.recommendationContext.typeahead.substring( + 0, + it.recommendationContext.typeahead.length - typeaheadChange.length + ) + } + it.recommendationContext.typeahead = newTypeahead } + } - for (i in 0..indexOfFirstNonWhiteSpace) { - val subTypeahead = typeahead.substring(i) - if (recommendation.startsWith(userInput + subTypeahead)) return subTypeahead + private fun updateSessionSelectedIndex(sessionContext: SessionContext) { + val selectedIndex = findNewSelectedIndex(false, sessionContext.selectedIndex) + if (selectedIndex == -1) { + LOG.debug { "None of the recommendation is valid at this point, cancelling the popup" } + CodeWhispererService.getInstance().disposeDisplaySession(false) + return } - return typeahead + + sessionContext.selectedIndex = selectedIndex } - fun updatePopupPanel(states: InvocationContext, sessionContext: SessionContext) { - val userInput = states.recommendationContext.userInputSinceInvocation - val details = states.recommendationContext.details + fun updatePopupPanel(sessionContext: SessionContext?) { + if (sessionContext == null || sessionContext.selectedIndex == -1 || sessionContext.isDisposed()) return val selectedIndex = sessionContext.selectedIndex - val typeaheadOriginal = sessionContext.typeaheadOriginal - val validCount = getValidCount(details, userInput, typeaheadOriginal) - val validSelectedIndex = getValidSelectedIndex(details, userInput, selectedIndex, typeaheadOriginal) + val previews = CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo() + if (selectedIndex >= previews.size) return + val validCount = getValidCount() + val validSelectedIndex = getValidSelectedIndex(selectedIndex) updateSelectedRecommendationLabelText(validSelectedIndex, validCount) updateNavigationPanel(validSelectedIndex, validCount) - updateImportPanel(details[selectedIndex].recommendation.mostRelevantMissingImports()) - updateCodeReferencePanel(states.requestContext.project, details[selectedIndex].recommendation.references()) + updateImportPanel(previews[selectedIndex].detail.recommendation.mostRelevantMissingImports()) + updateCodeReferencePanel(sessionContext.project, previews[selectedIndex].detail.recommendation.references()) } - fun render( - states: InvocationContext, - sessionContext: SessionContext, - overlappingLinesCount: Int, - isRecommendationAdded: Boolean, - isScrolling: Boolean, - ) { - updatePopupPanel(states, sessionContext) - - val caretPoint = states.requestContext.editor.offsetToXY(states.requestContext.caretPosition.offset) - sessionContext.seen.add(sessionContext.selectedIndex) + fun render(sessionContext: SessionContext, isRecommendationAdded: Boolean, isScrolling: Boolean) { + updatePopupPanel(sessionContext) // There are four cases that render() is called: // 1. Popup showing for the first time, both booleans are false, we should show the popup and update the latency @@ -232,20 +222,11 @@ class CodeWhispererPopupManager { // emit any events. // 4. User navigating through the completions or typing as the completion shows. We should not update the latency // end time and should not emit any events in this case. - if (!isRecommendationAdded) { - showPopup(states, sessionContext, states.popup, caretPoint, overlappingLinesCount) - if (!isScrolling) { - states.requestContext.latencyContext.codewhispererPostprocessingEnd = System.nanoTime() - states.requestContext.latencyContext.codewhispererEndToEndEnd = System.nanoTime() - } - } - if (isScrolling || - CodeWhispererInvocationStatus.getInstance().hasExistingInvocation() || - !sessionContext.isFirstTimeShowingPopup - ) { - return - } - CodeWhispererTelemetryService.getInstance().sendClientComponentLatencyEvent(states) + if (isRecommendationAdded) return + showPopup(sessionContext) + if (isScrolling) return + sessionContext.latencyContext.codewhispererPostprocessingEnd = System.nanoTime() + sessionContext.latencyContext.codewhispererEndToEndEnd = System.nanoTime() } fun dontClosePopupAndRun(runnable: () -> Unit) { @@ -257,84 +238,36 @@ class CodeWhispererPopupManager { } } - fun reset() { - sessionContext = SessionContext() - } - - fun cancelPopup(popup: JBPopup) { - popup.cancel() - Disposer.dispose(popup) - } - - fun closePopup(popup: JBPopup) { - popup.closeOk(null) - Disposer.dispose(popup) - } - - fun closePopup() { - myPopup?.let { - it.closeOk(null) - Disposer.dispose(it) + fun showPopup(sessionContext: SessionContext, force: Boolean = false) { + val p = sessionContext.editor.offsetToXY(sessionContext.popupOffset) + val popup: JBPopup? + if (sessionContext.popup == null) { + popup = initPopup() + sessionContext.popup = popup + CodeWhispererInvocationStatus.getInstance().setPopupStartTimestamp() + initPopupListener(sessionContext, popup) + } else { + popup = sessionContext.popup } - } - - fun showPopup( - states: InvocationContext, - sessionContext: SessionContext, - popup: JBPopup, - p: Point, - overlappingLinesCount: Int, - ) { - val editor = states.requestContext.editor - val detailContexts = states.recommendationContext.details - val userInputOriginal = states.recommendationContext.userInputOriginal - val userInput = states.recommendationContext.userInputSinceInvocation - val selectedIndex = sessionContext.selectedIndex - val typeaheadOriginal = sessionContext.typeaheadOriginal - val typeahead = sessionContext.typeahead + val editor = sessionContext.editor + val previews = CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo() + val userInputOriginal = previews[sessionContext.selectedIndex].userInput val userInputLines = userInputOriginal.split("\n").size - 1 - val lineCount = getReformattedRecommendation(detailContexts[selectedIndex], userInput).split("\n").size - val additionalLines = typeaheadOriginal.split("\n").size - typeahead.split("\n").size val popupSize = (popup as AbstractPopup).preferredContentSize - val yBelowLastLine = p.y + (lineCount + additionalLines + userInputLines - overlappingLinesCount) * editor.lineHeight - val yAboveFirstLine = p.y - popupSize.height + (additionalLines + userInputLines) * editor.lineHeight + val yAboveFirstLine = p.y - popupSize.height + userInputLines * editor.lineHeight + val popupRect = Rectangle(p.x, yAboveFirstLine, popupSize.width, popupSize.height) val editorRect = editor.scrollingModel.visibleArea - var popupRect = Rectangle(p.x, yBelowLastLine, popupSize.width, popupSize.height) var shouldHidePopup = false - CodeWhispererInvocationStatus.getInstance().setPopupActive(true) + CodeWhispererInvocationStatus.getInstance().setDisplaySessionActive(true) - // Check if the current editor still has focus. If not, don't show the popup. - val isSameEditorAsTrigger = if (!AppMode.isRemoteDevHost()) { - editor.contentComponent.isFocusOwner - } else { - FileEditorManager.getInstance(states.requestContext.project).selectedTextEditorWithRemotes.firstOrNull() == editor - } - if (!isSameEditorAsTrigger) { - LOG.debug { "Current editor no longer has focus, not showing the popup" } - cancelPopup(popup) - return + if (!editorRect.contains(popupRect)) { + // popup location above first line don't work, so don't show the popup + shouldHidePopup = true } - val popupLocation = - if (!editorRect.contains(popupRect)) { - popupRect = Rectangle(p.x, yAboveFirstLine, popupSize.width, popupSize.height) - if (!editorRect.contains(popupRect)) { - // both popup location (below last line and above first line) don't work, so don't show the popup - shouldHidePopup = true - } - LOG.debug { - "Show popup above the first line of recommendation. " + - "Editor position: $editorRect, popup position: $popupRect" - } - Point(p.x, yAboveFirstLine) - } else { - LOG.debug { - "Show popup below the last line of recommendation. " + - "Editor position: $editorRect, popup position: $popupRect" - } - Point(p.x, yBelowLastLine) - } + // popup to always display above the current editing line + val popupLocation = Point(p.x, yAboveFirstLine) val relativePopupLocationToEditor = RelativePoint(editor.contentComponent, popupLocation) @@ -347,8 +280,11 @@ class CodeWhispererPopupManager { } } else { if (!AppMode.isRemoteDevHost()) { - popup.show(relativePopupLocationToEditor) + if (force && !shouldHidePopup) { + popup.show(relativePopupLocationToEditor) + } } else { + // TODO: Fix in remote case the popup should display above the current editing line // TODO: For now, the popup will always display below the suggestions, without checking // if the location the popup is about to show at stays in the editor window or not, due to // the limitation of BackendBeAbstractPopup @@ -363,22 +299,6 @@ class CodeWhispererPopupManager { editor.putUserData(PopupFactoryImpl.ANCHOR_POPUP_POSITION, popupPositionForRemote) popup.showInBestPositionFor(editor) } - val perceivedLatency = CodeWhispererInvocationStatus.getInstance().getTimeSinceDocumentChanged() - CodeWhispererTelemetryService.getInstance().sendPerceivedLatencyEvent( - detailContexts[selectedIndex].requestId, - states.requestContext, - states.responseContext, - perceivedLatency - ) - } - - // popup.popupWindow is null in remote host - if (!AppMode.isRemoteDevHost()) { - if (shouldHidePopup) { - WindowManager.getInstance().setAlphaModeRatio(popup.popupWindow, 1f) - } else { - WindowManager.getInstance().setAlphaModeRatio(popup.popupWindow, 0.1f) - } } } @@ -386,70 +306,68 @@ class CodeWhispererPopupManager { .createComponentPopupBuilder(popupComponents.panel, null) .setAlpha(0.1F) .setCancelOnClickOutside(true) - .setCancelOnOtherWindowOpen(true) - .setCancelKeyEnabled(true) .setCancelOnWindowDeactivation(true) - .createPopup().also { - myPopup = it - } + .createPopup() fun getReformattedRecommendation(detailContext: DetailContext, userInput: String) = detailContext.reformatted.content().substring(userInput.length) - fun initPopupListener(states: InvocationContext) { - addPopupListener(states) - states.requestContext.editor.scrollingModel.addVisibleAreaListener(CodeWhispererScrollListener(states), states) - addButtonActionListeners(states) - addMessageSubscribers(states) - setPopupActionHandlers(states) - addComponentListeners(states) + private fun initPopupListener(sessionContext: SessionContext, popup: JBPopup) { + addPopupListener(popup) + sessionContext.editor.scrollingModel.addVisibleAreaListener(CodeWhispererScrollListener(sessionContext), sessionContext) + addButtonActionListeners(sessionContext) + addMessageSubscribers(sessionContext) + setPopupActionHandlers(sessionContext) + addComponentListeners(sessionContext) } - private fun addPopupListener(states: InvocationContext) { - val listener = CodeWhispererPopupListener(states) - states.popup.addListener(listener) - Disposer.register(states) { states.popup.removeListener(listener) } + private fun addPopupListener(popup: JBPopup) { + val listener = CodeWhispererPopupListener() + popup.addListener(listener) + Disposer.register(popup) { + popup.removeListener(listener) + } } - private fun addMessageSubscribers(states: InvocationContext) { - val connect = ApplicationManager.getApplication().messageBus.connect(states) + private fun addMessageSubscribers(sessionContext: SessionContext) { + val connect = ApplicationManager.getApplication().messageBus.connect(sessionContext) connect.subscribe( CODEWHISPERER_USER_ACTION_PERFORMED, object : CodeWhispererUserActionListener { - override fun navigateNext(states: InvocationContext) { - changeStates(states, 1, "", true) + override fun navigateNext(sessionContext: SessionContext) { + changeStatesForNavigation(sessionContext, 1) } - override fun navigatePrevious(states: InvocationContext) { - changeStates(states, -1, "", true) + override fun navigatePrevious(sessionContext: SessionContext) { + changeStatesForNavigation(sessionContext, -1) } - override fun backspace(states: InvocationContext, diff: String) { - changeStates(states, 0, diff, false) + override fun backspace(sessionContext: SessionContext, diff: String) { + changeStatesForTypeahead(sessionContext, diff, false) } - override fun enter(states: InvocationContext, diff: String) { - changeStates(states, 0, diff, true) + override fun enter(sessionContext: SessionContext, diff: String) { + changeStatesForTypeahead(sessionContext, diff, true) } - override fun type(states: InvocationContext, diff: String) { + override fun type(sessionContext: SessionContext, diff: String) { // remove the character at primaryCaret if it's the same as the typed character - val caretOffset = states.requestContext.editor.caretModel.primaryCaret.offset - val document = states.requestContext.editor.document + val caretOffset = sessionContext.editor.caretModel.primaryCaret.offset + val document = sessionContext.editor.document val text = document.charsSequence if (caretOffset < text.length && diff == text[caretOffset].toString()) { - WriteCommandAction.runWriteCommandAction(states.requestContext.project) { + WriteCommandAction.runWriteCommandAction(sessionContext.project) { document.deleteString(caretOffset, caretOffset + 1) } } - changeStates(states, 0, diff, true) + changeStatesForTypeahead(sessionContext, diff, true) } - override fun beforeAccept(states: InvocationContext, sessionContext: SessionContext) { + override fun beforeAccept(sessionContext: SessionContext) { dontClosePopupAndRun { - CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext) + CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(sessionContext) } - closePopup(states.popup) + CodeWhispererService.getInstance().disposeDisplaySession(true) } } ) @@ -492,60 +410,75 @@ class CodeWhispererPopupManager { Disposer.register(newHandler.states) { EditorActionManager.getInstance().setActionHandler(id, oldHandler) } } - private fun addComponentListeners(states: InvocationContext) { - val editor = states.requestContext.editor - val codewhispererSelectionListener: SelectionListener = object : SelectionListener { + private fun addComponentListeners(sessionContext: SessionContext) { + val editor = sessionContext.editor + val codeWhispererSelectionListener: SelectionListener = object : SelectionListener { override fun selectionChanged(event: SelectionEvent) { if (shouldListenerCancelPopup) { - cancelPopup(states.popup) + CodeWhispererService.getInstance().disposeDisplaySession(false) } super.selectionChanged(event) } } - editor.selectionModel.addSelectionListener(codewhispererSelectionListener) - Disposer.register(states) { editor.selectionModel.removeSelectionListener(codewhispererSelectionListener) } + editor.selectionModel.addSelectionListener(codeWhispererSelectionListener) + Disposer.register(sessionContext) { editor.selectionModel.removeSelectionListener(codeWhispererSelectionListener) } - val codewhispererDocumentListener: DocumentListener = object : DocumentListener { + val codeWhispererDocumentListener: DocumentListener = object : DocumentListener { override fun documentChanged(event: DocumentEvent) { if (shouldListenerCancelPopup) { - cancelPopup(states.popup) + // handle IntelliSense accept case + // TODO: handle bulk delete (delete word) case + if (editor.document == event.document && + editor.caretModel.offset == event.offset && + event.newLength > event.oldLength + ) { + dontClosePopupAndRun { + super.documentChanged(event) + editor.caretModel.moveCaretRelatively(event.newLength, 0, false, false, true) + changeStatesForTypeahead(sessionContext, event.newFragment.toString(), true) + } + return + } else { + CodeWhispererService.getInstance().disposeDisplaySession(false) + } } super.documentChanged(event) } } - editor.document.addDocumentListener(codewhispererDocumentListener, states) + editor.document.addDocumentListener(codeWhispererDocumentListener, sessionContext) - val codewhispererCaretListener: CaretListener = object : CaretListener { + val codeWhispererCaretListener: CaretListener = object : CaretListener { override fun caretPositionChanged(event: CaretEvent) { if (shouldListenerCancelPopup) { - cancelPopup(states.popup) + CodeWhispererService.getInstance().disposeDisplaySession(false) } super.caretPositionChanged(event) } } - editor.caretModel.addCaretListener(codewhispererCaretListener) - Disposer.register(states) { editor.caretModel.removeCaretListener(codewhispererCaretListener) } + editor.caretModel.addCaretListener(codeWhispererCaretListener) + Disposer.register(sessionContext) { editor.caretModel.removeCaretListener(codeWhispererCaretListener) } val editorComponent = editor.contentComponent if (editorComponent.isShowing) { val window = ComponentUtil.getWindow(editorComponent) val windowListener: ComponentListener = object : ComponentAdapter() { - override fun componentMoved(event: ComponentEvent) { - cancelPopup(states.popup) + override fun componentMoved(e: ComponentEvent) { + CodeWhispererService.getInstance().disposeDisplaySession(false) + super.componentMoved(e) } override fun componentShown(e: ComponentEvent?) { - cancelPopup(states.popup) + CodeWhispererService.getInstance().disposeDisplaySession(false) super.componentShown(e) } } window?.addComponentListener(windowListener) - Disposer.register(states) { window?.removeComponentListener(windowListener) } + Disposer.register(sessionContext) { window?.removeComponentListener(windowListener) } } } private fun updateSelectedRecommendationLabelText(validSelectedIndex: Int, validCount: Int) { - if (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) { + if (CodeWhispererInvocationStatus.getInstance().hasExistingServiceInvocation()) { popupComponents.recommendationInfoLabel.text = message("codewhisperer.popup.pagination_info") LOG.debug { "Pagination in progress. Current total: $validCount" } } else { @@ -622,14 +555,10 @@ class CodeWhispererPopupManager { ParameterInfoController.existsWithVisibleHintForEditor(editor, true) || LookupManager.getActiveLookup(editor) != null - private fun findNewSelectedIndex( - isReverse: Boolean, - detailContexts: List, - userInput: String, - start: Int, - typeahead: String, - ): Int { - val count = detailContexts.size + fun findNewSelectedIndex(isReverse: Boolean, selectedIndex: Int): Int { + val start = if (selectedIndex == -1) 0 else selectedIndex + val previews = CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo() + val count = previews.size val unit = if (isReverse) -1 else 1 var currIndex: Int for (i in 0 until count) { @@ -637,45 +566,34 @@ class CodeWhispererPopupManager { if (currIndex < 0) { currIndex += count } - if (isValidRecommendation(detailContexts[currIndex], userInput, typeahead)) { + if (isValidRecommendation(previews[currIndex])) { return currIndex } } return -1 } - private fun getValidCount(detailContexts: List, userInput: String, typeahead: String): Int = - detailContexts.filter { isValidRecommendation(it, userInput, typeahead) }.size - - private fun getValidSelectedIndex( - detailContexts: List, - userInput: String, - selectedIndex: Int, - typeahead: String, - ): Int { - var currIndexIgnoreInvalid = 0 - detailContexts.forEachIndexed { index, value -> + private fun getValidCount(): Int = + CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo().filter { isValidRecommendation(it) }.size + + private fun getValidSelectedIndex(selectedIndex: Int): Int { + var curr = 0 + + val previews = CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo() + previews.forEachIndexed { index, preview -> if (index == selectedIndex) { - return currIndexIgnoreInvalid + return curr } - if (isValidRecommendation(value, userInput, typeahead)) { - currIndexIgnoreInvalid++ + if (isValidRecommendation(preview)) { + curr++ } } return -1 } - private fun isValidRecommendation(detailContext: DetailContext, userInput: String, typeahead: String): Boolean { - if (detailContext.isDiscarded) return false - if (detailContext.recommendation.content().isEmpty()) return false - val indexOfFirstNonWhiteSpace = typeahead.indexOfFirst { !it.isWhitespace() } - if (indexOfFirstNonWhiteSpace == -1) return true - - for (i in 0..indexOfFirstNonWhiteSpace) { - val subTypeahead = typeahead.substring(i) - if (detailContext.reformatted.content().startsWith(userInput + subTypeahead)) return true - } - return false + private fun isValidRecommendation(preview: PreviewContext): Boolean { + if (preview.detail.isDiscarded) return false + return preview.detail.recommendation.content().startsWith(preview.userInput + preview.typeahead) } companion object { @@ -693,17 +611,17 @@ class CodeWhispererPopupManager { } interface CodeWhispererPopupStateChangeListener { - fun stateChanged(states: InvocationContext, sessionContext: SessionContext) {} - fun scrolled(states: InvocationContext, sessionContext: SessionContext) {} + fun stateChanged(sessionContext: SessionContext) {} + fun scrolled(sessionContext: SessionContext) {} fun recommendationAdded(states: InvocationContext, sessionContext: SessionContext) {} } interface CodeWhispererUserActionListener { - fun backspace(states: InvocationContext, diff: String) {} - fun enter(states: InvocationContext, diff: String) {} - fun type(states: InvocationContext, diff: String) {} - fun navigatePrevious(states: InvocationContext) {} - fun navigateNext(states: InvocationContext) {} - fun beforeAccept(states: InvocationContext, sessionContext: SessionContext) {} - fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) {} + fun backspace(sessionContext: SessionContext, diff: String) {} + fun enter(sessionContext: SessionContext, diff: String) {} + fun type(sessionContext: SessionContext, diff: String) {} + fun navigatePrevious(sessionContext: SessionContext) {} + fun navigateNext(sessionContext: SessionContext) {} + fun beforeAccept(sessionContext: SessionContext) {} + fun afterAccept(states: InvocationContext, previews: List, sessionContext: SessionContext, rangeMarker: RangeMarker) {} } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererUIChangeListener.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererUIChangeListener.kt index e4bb87feec8..e5dce240995 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererUIChangeListener.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererUIChangeListener.kt @@ -15,22 +15,26 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationCo import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationChunk import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererRecommendationManager +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService class CodeWhispererUIChangeListener : CodeWhispererPopupStateChangeListener { - override fun stateChanged(states: InvocationContext, sessionContext: SessionContext) { - val editor = states.requestContext.editor + override fun stateChanged(sessionContext: SessionContext) { + val editor = sessionContext.editor val editorManager = CodeWhispererEditorManager.getInstance() + val previews = CodeWhispererService.getInstance().getAllSuggestionsPreviewInfo() val selectedIndex = sessionContext.selectedIndex - val typeahead = sessionContext.typeahead - val detail = states.recommendationContext.details[selectedIndex] + val typeahead = previews[selectedIndex].typeahead + val detail = previews[selectedIndex].detail val caretOffset = editor.caretModel.primaryCaret.offset val document = editor.document val lineEndOffset = document.getLineEndOffset(document.getLineNumber(caretOffset)) + detail.hasSeen = true + // get matching brackets from recommendations to the brackets after caret position val remaining = CodeWhispererPopupManager.getInstance().getReformattedRecommendation( detail, - states.recommendationContext.userInputSinceInvocation + previews[selectedIndex].userInput, ).substring(typeahead.length) val remainingLines = remaining.split("\n") @@ -61,7 +65,7 @@ class CodeWhispererUIChangeListener : CodeWhispererPopupStateChangeListener { }, HighlighterTargetArea.EXACT_RANGE ) - Disposer.register(states.popup) { + Disposer.register(sessionContext) { editor.markupModel.removeHighlighter(rangeHighlighter) } sessionContext.toBeRemovedHighlighter = rangeHighlighter @@ -87,57 +91,21 @@ class CodeWhispererUIChangeListener : CodeWhispererPopupStateChangeListener { // inlay chunks are chunks from first line(chunks) and an additional chunk from other lines val inlayChunks = chunks + listOf(RecommendationChunk(otherLinesInlayText, 0, chunks.last().inlayOffset)) - CodeWhispererInlayManager.getInstance().updateInlays(states, inlayChunks) + CodeWhispererInlayManager.getInstance().updateInlays(sessionContext, inlayChunks) CodeWhispererPopupManager.getInstance().render( - states, sessionContext, - overlappingLinesCount, isRecommendationAdded = false, isScrolling = false ) } - override fun scrolled(states: InvocationContext, sessionContext: SessionContext) { - if (states.popup.isDisposed) return - val editor = states.requestContext.editor - val editorManager = CodeWhispererEditorManager.getInstance() - val selectedIndex = sessionContext.selectedIndex - val typeahead = sessionContext.typeahead - val detail = states.recommendationContext.details[selectedIndex] - - // get matching brackets from recommendations to the brackets after caret position - val remaining = CodeWhispererPopupManager.getInstance().getReformattedRecommendation( - detail, - states.recommendationContext.userInputSinceInvocation - ).substring(typeahead.length) - - val remainingLines = remaining.split("\n") - val otherLinesOfRemaining = remainingLines.drop(1) - - // process other lines inlays, where we do tail-head matching as much as possible - val overlappingLinesCount = editorManager.findOverLappingLines( - editor, - otherLinesOfRemaining, - detail.isTruncatedOnRight, - sessionContext - ) - - CodeWhispererPopupManager.getInstance().render( - states, - sessionContext, - overlappingLinesCount, - isRecommendationAdded = false, - isScrolling = true - ) + override fun scrolled(sessionContext: SessionContext) { + sessionContext.isFirstTimeShowingPopup = false + CodeWhispererPopupManager.getInstance().render(sessionContext, isRecommendationAdded = false, isScrolling = true) } override fun recommendationAdded(states: InvocationContext, sessionContext: SessionContext) { - CodeWhispererPopupManager.getInstance().render( - states, - sessionContext, - 0, - isRecommendationAdded = true, - isScrolling = false - ) + sessionContext.isFirstTimeShowingPopup = false + CodeWhispererPopupManager.getInstance().render(sessionContext, isRecommendationAdded = true, isScrolling = false) } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererInvocationStatus.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererInvocationStatus.kt index 2a2d70bd903..4c58139b354 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererInvocationStatus.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererInvocationStatus.kt @@ -16,9 +16,8 @@ import java.util.concurrent.atomic.AtomicBoolean @Service class CodeWhispererInvocationStatus { - private val isInvokingCodeWhisperer: AtomicBoolean = AtomicBoolean(false) + private val isInvokingService: AtomicBoolean = AtomicBoolean(false) private var invokingSessionId: String? = null - private var timeAtLastInvocationComplete: Instant? = null var timeAtLastDocumentChanged: Instant = Instant.now() private set private var isPopupActive: Boolean = false @@ -26,30 +25,21 @@ class CodeWhispererInvocationStatus { var popupStartTimestamp: Instant? = null private set - fun checkExistingInvocationAndSet(): Boolean = - if (isInvokingCodeWhisperer.getAndSet(true)) { - LOG.debug { "Have existing CodeWhisperer invocation, sessionId: $invokingSessionId" } - true - } else { - ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(true) - LOG.debug { "Starting CodeWhisperer invocation" } - false - } + fun startInvocation() { + isInvokingService.set(true) + LOG.debug { "Starting CodeWhisperer invocation" } + } - fun hasExistingInvocation(): Boolean = isInvokingCodeWhisperer.get() + fun hasExistingServiceInvocation(): Boolean = isInvokingService.get() fun finishInvocation() { - if (isInvokingCodeWhisperer.compareAndSet(true, false)) { + if (isInvokingService.compareAndSet(true, false)) { ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(false) LOG.debug { "Ending CodeWhisperer invocation" } invokingSessionId = null } } - fun setInvocationComplete() { - timeAtLastInvocationComplete = Instant.now() - } - fun documentChanged() { timeAtLastDocumentChanged = Instant.now() } @@ -69,9 +59,9 @@ class CodeWhispererInvocationStatus { return timeCanShowCodeWhisperer.isBefore(Instant.now()) } - fun isPopupActive(): Boolean = isPopupActive + fun isDisplaySessionActive(): Boolean = isPopupActive - fun setPopupActive(value: Boolean) { + fun setDisplaySessionActive(value: Boolean) { isPopupActive = value } @@ -84,11 +74,6 @@ class CodeWhispererInvocationStatus { invokingSessionId = sessionId } - fun hasEnoughDelayToInvokeCodeWhisperer(): Boolean { - val timeCanShowCodeWhisperer = timeAtLastInvocationStart?.plusMillis(CodeWhispererConstants.INVOCATION_INTERVAL) ?: return true - return timeCanShowCodeWhisperer.isBefore(Instant.now()) - } - companion object { private val LOG = getLogger() fun getInstance(): CodeWhispererInvocationStatus = service() diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt index 2099b963999..7685e0a5c19 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt @@ -9,6 +9,7 @@ import org.jetbrains.annotations.VisibleForTesting import software.amazon.awssdk.services.codewhispererruntime.model.Completion import software.amazon.awssdk.services.codewhispererruntime.model.Span import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationChunk import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType import kotlin.math.max @@ -63,7 +64,7 @@ class CodeWhispererRecommendationManager { userInput: String, recommendations: List, requestId: String, - ): List { + ): MutableList { val seen = mutableSetOf() return recommendations.map { val isDiscardedByUserInput = !it.content().startsWith(userInput) || it.content() == userInput @@ -126,7 +127,7 @@ class CodeWhispererRecommendationManager { overlap, getCompletionType(it) ) - } + }.toMutableList() } fun findRightContextOverlap( diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt index 8b9c9d22949..80d01722c57 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt @@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.service import com.intellij.codeInsight.hint.HintManager import com.intellij.notification.NotificationAction import com.intellij.openapi.Disposable +import com.intellij.openapi.actionSystem.DataKey import com.intellij.openapi.application.ApplicationInfo import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.application.runInEdt @@ -15,7 +16,6 @@ import com.intellij.openapi.components.service import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.VisualPosition import com.intellij.openapi.project.Project -import com.intellij.openapi.ui.popup.JBPopup import com.intellij.openapi.util.Disposer import com.intellij.psi.PsiDocumentManager import com.intellij.psi.PsiFile @@ -25,6 +25,7 @@ import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Deferred import kotlinx.coroutines.Job import kotlinx.coroutines.async +import kotlinx.coroutines.cancel import kotlinx.coroutines.delay import kotlinx.coroutines.isActive import kotlinx.coroutines.launch @@ -46,7 +47,6 @@ import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection @@ -64,6 +64,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContex import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo @@ -88,12 +89,17 @@ import software.aws.toolkits.resources.message import software.aws.toolkits.telemetry.CodewhispererCompletionType import software.aws.toolkits.telemetry.CodewhispererSuggestionState import software.aws.toolkits.telemetry.CodewhispererTriggerType +import java.util.concurrent.CancellationException import java.util.concurrent.TimeUnit @Service class CodeWhispererService(private val cs: CoroutineScope) : Disposable { private val codeInsightSettingsFacade = CodeInsightsSettingsFacade() private var refreshFailure: Int = 0 + private val ongoingRequests = mutableMapOf() + val ongoingRequestsContext = mutableMapOf() + private var jobId = 0 + private var sessionContext: SessionContext? = null init { Disposer.register(this, codeInsightSettingsFacade) @@ -161,18 +167,27 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { val isInjectedFile = runReadAction { psiFile.isInjectedText() } if (isInjectedFile) return + val currentJobId = jobId++ val requestContext = try { - getRequestContext(triggerTypeInfo, editor, project, psiFile, latencyContext) + getRequestContext(triggerTypeInfo, editor, project, psiFile) } catch (e: Exception) { LOG.debug { e.message.toString() } CodeWhispererTelemetryService.getInstance().sendFailedServiceInvocationEvent(project, e::class.simpleName) return } + val caretContext = requestContext.fileContextInfo.caretContext + ongoingRequestsContext.forEach { (k, v) -> + val vCaretContext = v.fileContextInfo.caretContext + if (vCaretContext == caretContext) { + LOG.debug { "same caretContext found from job: $k, left context ${vCaretContext.leftContextOnCurrentLine}, jobId: $currentJobId" } + return + } + } val language = requestContext.fileContextInfo.programmingLanguage val leftContext = requestContext.fileContextInfo.caretContext.leftFileContext if (!language.isCodeCompletionSupported() || (checkLeftContextKeywordsForJsonAndYaml(leftContext, language.languageId))) { - LOG.debug { "Programming language $language is not supported by CodeWhisperer" } + LOG.debug { "Programming language $language is not supported by CodeWhisperer, jobId: $currentJobId" } if (triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) { showCodeWhispererInfoHint( requestContext.editor, @@ -183,7 +198,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { } LOG.debug { - "Calling CodeWhisperer service, trigger type: ${triggerTypeInfo.triggerType}" + + "Calling CodeWhisperer service, jobId: $currentJobId, trigger type: ${triggerTypeInfo.triggerType}" + if (triggerTypeInfo.triggerType == CodewhispererTriggerType.AutoTrigger) { ", auto-trigger type: ${triggerTypeInfo.automatedTriggerType}" } else { @@ -191,29 +206,28 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { } } - val invocationStatus = CodeWhispererInvocationStatus.getInstance() - if (invocationStatus.checkExistingInvocationAndSet()) { - return - } + CodeWhispererInvocationStatus.getInstance().startInvocation() - invokeCodeWhispererInBackground(requestContext) + invokeCodeWhispererInBackground(requestContext, currentJobId, latencyContext) } - internal fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job { - val popup = CodeWhispererPopupManager.getInstance().initPopup() - Disposer.register(popup) { CodeWhispererInvocationStatus.getInstance().finishInvocation() } + internal fun invokeCodeWhispererInBackground(requestContext: RequestContext, currentJobId: Int, latencyContext: LatencyContext): Job? { + ongoingRequestsContext[currentJobId] = requestContext + val sessionContext = sessionContext ?: SessionContext(requestContext.project, requestContext.editor, latencyContext = latencyContext) + + // In rare cases when there's an ongoing session and subsequent triggers are from a different project or editor -- + // we will cancel the existing session(since we've already moved to a different project or editor simply return. + if (requestContext.project != sessionContext.project || requestContext.editor != sessionContext.editor) { + disposeDisplaySession(false) + return null + } + this.sessionContext = sessionContext val workerContexts = mutableListOf() - // When popup is disposed we will cancel this coroutine. The only places popup can get disposed should be - // from CodeWhispererPopupManager.cancelPopup() and CodeWhispererPopupManager.closePopup(). - // It's possible and ok that coroutine will keep running until the next time we check it's state. - // As long as we don't show to the user extra info we are good. - val coroutineScope = disposableCoroutineScope(popup) - var states: InvocationContext? = null var lastRecommendationIndex = -1 - val job = coroutineScope.launch { + val job = cs.launch { try { val responseIterable = CodeWhispererClientAdaptor.getInstance(requestContext.project).generateCompletionsPaginator( buildCodeWhispererRequest( @@ -224,8 +238,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { ) var startTime = System.nanoTime() - requestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime() - requestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime() + latencyContext.codewhispererPreprocessingEnd = System.nanoTime() + latencyContext.paginationAllCompletionsStart = System.nanoTime() CodeWhispererInvocationStatus.getInstance().setInvocationStart() var requestCount = 0 for (response in responseIterable) { @@ -236,13 +250,13 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { val requestId = response.responseMetadata().requestId() val sessionId = response.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0] if (requestCount == 1) { - requestContext.latencyContext.codewhispererPostprocessingStart = System.nanoTime() - requestContext.latencyContext.paginationFirstCompletionTime = latency - requestContext.latencyContext.firstRequestId = requestId + latencyContext.codewhispererPostprocessingStart = System.nanoTime() + latencyContext.paginationFirstCompletionTime = latency + latencyContext.firstRequestId = requestId CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId) } if (response.nextToken().isEmpty()) { - requestContext.latencyContext.paginationAllCompletionsEnd = System.nanoTime() + latencyContext.paginationAllCompletionsEnd = System.nanoTime() } val responseContext = ResponseContext(sessionId) logServiceInvocation(requestId, requestContext, responseContext, response.completions(), latency, null) @@ -250,6 +264,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_CODE_COMPLETION_PERFORMED) .onSuccess(requestContext.fileContextInfo) CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent( + currentJobId, requestId, requestContext, responseContext, @@ -271,11 +286,11 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { // task hasn't been finished yet, in this case simply add another task to the queue. If they // see worker queue is empty, the previous tasks must have been finished before this. In this // case render CodeWhisperer UI directly. - val workerContext = WorkerContext(requestContext, responseContext, validatedResponse, popup) + val workerContext = WorkerContext(requestContext, responseContext, validatedResponse) if (workerContexts.isNotEmpty()) { workerContexts.add(workerContext) } else { - if (states == null && !popup.isDisposed && + if (ongoingRequests.values.filterNotNull().isEmpty() && !CodeWhispererInvocationStatus.getInstance().hasEnoughDelayToShowCodeWhisperer() ) { // It's the first response, and no enough delay before showing @@ -285,7 +300,16 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { } runInEdt { workerContexts.forEach { - states = processCodeWhispererUI(it, states) + processCodeWhispererUI( + sessionContext, + it, + ongoingRequests[currentJobId], + cs, + currentJobId + ) + if (!ongoingRequests.contains(currentJobId)) { + cs.coroutineContext.cancel() + } } workerContexts.clear() } @@ -293,14 +317,25 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { workerContexts.add(workerContext) } else { // Have enough delay before showing for the first response, or it's subsequent responses - states = processCodeWhispererUI(workerContext, states) + processCodeWhispererUI( + sessionContext, + workerContext, + ongoingRequests[currentJobId], + cs, + currentJobId + ) } } } if (!isActive) { // If job is cancelled before we do another request, don't bother making // another API call to save resources - LOG.debug { "Skipping sending remaining requests on CodeWhisperer session exit" } + LOG.debug { "Skipping sending remaining requests on inactive CodeWhisperer session exit" } + return@launch + } + if (requestCount >= PAGINATION_REQUEST_COUNT_ALLOWED) { + LOG.debug { "Only $PAGINATION_REQUEST_COUNT_ALLOWED request per pagination session for now" } + CodeWhispererInvocationStatus.getInstance().finishInvocation() break } } @@ -321,6 +356,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { val responseContext = ResponseContext(sessionId) CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent( + currentJobId, requestId, requestContext, responseContext, @@ -350,7 +386,6 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { ) ) CodeWhispererInvocationStatus.getInstance().finishInvocation() - CodeWhispererInvocationStatus.getInstance().setInvocationComplete() requestContext.customizationArn?.let { CodeWhispererModelConfigurator.getInstance().invalidateCustomization(it) } @@ -358,7 +393,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { showRecommendationsInPopup( requestContext.editor, requestContext.triggerTypeInfo, - requestContext.latencyContext + latencyContext ) } return@launch @@ -366,7 +401,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { requestId = e.requestId() ?: "" sessionId = e.awsErrorDetails().sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0] displayMessage = e.awsErrorDetails().errorMessage() ?: message("codewhisperer.trigger.error.server_side") - } else if (e is software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException) { + } else if (e is CodeWhispererRuntimeException) { requestId = e.requestId() ?: "" sessionId = e.awsErrorDetails().sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0] displayMessage = e.awsErrorDetails().errorMessage() ?: message("codewhisperer.trigger.error.server_side") @@ -389,6 +424,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId) logServiceInvocation(requestId, requestContext, responseContext, emptyList(), null, exceptionType) CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent( + currentJobId, requestId, requestContext, responseContext, @@ -410,7 +446,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { // We should only show error hint when CodeWhisperer popup is not visible, // and make it silent if CodeWhisperer popup is showing. runInEdt { - if (!CodeWhispererInvocationStatus.getInstance().isPopupActive()) { + if (!CodeWhispererInvocationStatus.getInstance().isDisplaySessionActive()) { showCodeWhispererErrorHint(requestContext.editor, displayMessage) } } @@ -418,15 +454,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { } CodeWhispererInvocationStatus.getInstance().finishInvocation() runInEdt { - states?.let { - CodeWhispererPopupManager.getInstance().updatePopupPanel( - it, - CodeWhispererPopupManager.getInstance().sessionContext - ) - } + CodeWhispererPopupManager.getInstance().updatePopupPanel(sessionContext) } - } finally { - CodeWhispererInvocationStatus.getInstance().setInvocationComplete() } } @@ -434,61 +463,66 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { } @RequiresEdt - private fun processCodeWhispererUI(workerContext: WorkerContext, currStates: InvocationContext?): InvocationContext? { + private fun processCodeWhispererUI( + sessionContext: SessionContext, + workerContext: WorkerContext, + currStates: InvocationContext?, + coroutine: CoroutineScope, + jobId: Int + ) { val requestContext = workerContext.requestContext val responseContext = workerContext.responseContext val response = workerContext.response - val popup = workerContext.popup val requestId = response.responseMetadata().requestId() // At this point when we are in EDT, the state of the popup will be thread-safe // across this thread execution, so if popup is disposed, we will stop here. // This extra check is needed because there's a time between when we get the response and // when we enter the EDT. - if (popup.isDisposed) { - LOG.debug { "Stop showing CodeWhisperer recommendations on CodeWhisperer session exit. RequestId: $requestId" } - return null + if (!coroutine.isActive || sessionContext.isDisposed()) { + LOG.debug { "Stop showing CodeWhisperer recommendations on CodeWhisperer session exit. RequestId: $requestId, jobId: $jobId" } + return } if (requestContext.editor.isDisposed) { - LOG.debug { "Stop showing CodeWhisperer recommendations since editor is disposed. RequestId: $requestId" } - CodeWhispererPopupManager.getInstance().cancelPopup(popup) - return null + LOG.debug { "Stop showing all CodeWhisperer recommendations since editor is disposed. RequestId: $requestId, jobId: $jobId" } + disposeDisplaySession(false) + return } - if (response.nextToken().isEmpty()) { - CodeWhispererInvocationStatus.getInstance().finishInvocation() - } + CodeWhispererInvocationStatus.getInstance().finishInvocation() val caretMovement = CodeWhispererEditorManager.getInstance().getCaretMovement( requestContext.editor, requestContext.caretPosition ) - val isPopupShowing: Boolean + val isPopupShowing = checkRecommendationsValidity(currStates, false) val nextStates: InvocationContext? if (currStates == null) { - // first response - nextStates = initStates(requestContext, responseContext, response, caretMovement, popup) - isPopupShowing = false + // first response for the jobId + nextStates = initStates(jobId, requestContext, responseContext, response, caretMovement, coroutine) - // receiving a null state means caret has moved backward or there's a conflict with - // Intellisense popup, so we are going to cancel the job + // receiving a null state means caret has moved backward, + // so we are going to cancel the current job if (nextStates == null) { - LOG.debug { "Cancelling popup and exiting CodeWhisperer session. RequestId: $requestId" } - CodeWhispererPopupManager.getInstance().cancelPopup(popup) - return null + return } } else { - // subsequent responses + // subsequent responses for the jobId nextStates = updateStates(currStates, response) - isPopupShowing = checkRecommendationsValidity(currStates, false) } + LOG.debug { "Adding ${response.completions().size} completions to the session. RequestId: $requestId, jobId: $jobId" } - val hasAtLeastOneValid = checkRecommendationsValidity(nextStates, response.nextToken().isEmpty()) + // TODO: may have bug when it's a mix of auto-trigger + manual trigger + val hasAtLeastOneValid = checkRecommendationsValidity(nextStates, true) + val allSuggestions = ongoingRequests.values.filterNotNull().flatMap { it.recommendationContext.details } + val valid = allSuggestions.filter { !it.isDiscarded }.size + LOG.debug { "Suggestions status: valid: $valid, discarded: ${allSuggestions.size - valid}" } // If there are no recommendations at all in this session, we need to manually send the user decision event here // since it won't be sent automatically later - if (nextStates.recommendationContext.details.isEmpty() && response.nextToken().isEmpty()) { + // TODO: may have bug; visit later + if (nextStates.recommendationContext.details.isEmpty()) { LOG.debug { "Received just an empty list from this session, requestId: $requestId" } CodeWhispererTelemetryService.getInstance().sendUserDecisionEvent( requestContext, @@ -508,38 +542,42 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { ) } if (!hasAtLeastOneValid) { - if (response.nextToken().isEmpty()) { - LOG.debug { "None of the recommendations are valid, exiting CodeWhisperer session" } - CodeWhispererPopupManager.getInstance().cancelPopup(popup) - return null + LOG.debug { "None of the recommendations are valid, exiting current CodeWhisperer pagination session" } + // If there's only one ongoing request, after disposing this, the entire session will also end + if (ongoingRequests.keys.size == 1) { + disposeDisplaySession(false) + } else { + disposeJob(jobId) + sessionContext.selectedIndex = CodeWhispererPopupManager.getInstance().findNewSelectedIndex(true, sessionContext.selectedIndex) } } else { - updateCodeWhisperer(nextStates, isPopupShowing) + updateCodeWhisperer(sessionContext, nextStates, isPopupShowing) } - return nextStates } private fun initStates( + jobId: Int, requestContext: RequestContext, responseContext: ResponseContext, response: GenerateCompletionsResponse, caretMovement: CaretMovement, - popup: JBPopup, + coroutine: CoroutineScope, ): InvocationContext? { val requestId = response.responseMetadata().requestId() val recommendations = response.completions() val visualPosition = requestContext.editor.caretModel.visualPosition - if (CodeWhispererPopupManager.getInstance().hasConflictingPopups(requestContext.editor)) { - LOG.debug { "Detect conflicting popup window with CodeWhisperer popup, not showing CodeWhisperer popup" } - sendDiscardedUserDecisionEventForAll(requestContext, responseContext, recommendations) - return null - } if (caretMovement == CaretMovement.MOVE_BACKWARD) { - LOG.debug { "Caret moved backward, discarding all of the recommendations. Request ID: $requestId" } - sendDiscardedUserDecisionEventForAll(requestContext, responseContext, recommendations) + LOG.debug { "Caret moved backward, discarding all of the recommendations and exiting the session. Request ID: $requestId, jobId: $jobId" } + val detailContexts = recommendations.map { + DetailContext("", it, it, true, false, "", getCompletionType(it)) + }.toMutableList() + val recommendationContext = RecommendationContext(detailContexts, "", "", VisualPosition(0, 0), jobId) + ongoingRequests[jobId] = buildInvocationContext(requestContext, responseContext, recommendationContext, coroutine) + disposeDisplaySession(false) return null } + val userInputOriginal = CodeWhispererEditorManager.getInstance().getUserInputSinceInvocation( requestContext.editor, requestContext.caretPosition.offset @@ -563,8 +601,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { recommendations, requestId ) - val recommendationContext = RecommendationContext(detailContexts, userInputOriginal, userInput, visualPosition) - return buildInvocationContext(requestContext, responseContext, recommendationContext, popup) + val recommendationContext = RecommendationContext(detailContexts, userInputOriginal, userInput, visualPosition, jobId) + ongoingRequests[jobId] = buildInvocationContext(requestContext, responseContext, recommendationContext, coroutine) + return ongoingRequests[jobId] } private fun updateStates( @@ -572,24 +611,19 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { response: GenerateCompletionsResponse, ): InvocationContext { val recommendationContext = states.recommendationContext - val details = recommendationContext.details val newDetailContexts = CodeWhispererRecommendationManager.getInstance().buildDetailContext( states.requestContext, recommendationContext.userInputSinceInvocation, response.completions(), response.responseMetadata().requestId() ) - Disposer.dispose(states) - val updatedStates = states.copy( - recommendationContext = recommendationContext.copy(details = details + newDetailContexts) - ) - Disposer.register(states.popup, updatedStates) - CodeWhispererPopupManager.getInstance().initPopupListener(updatedStates) - return updatedStates + recommendationContext.details.addAll(newDetailContexts) + return states } - private fun checkRecommendationsValidity(states: InvocationContext, showHint: Boolean): Boolean { + private fun checkRecommendationsValidity(states: InvocationContext?, showHint: Boolean): Boolean { + if (states == null) return false val details = states.recommendationContext.details // set to true when at least one is not discarded or empty @@ -604,35 +638,48 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { return hasAtLeastOneValid } - private fun updateCodeWhisperer(states: InvocationContext, recommendationAdded: Boolean) { - CodeWhispererPopupManager.getInstance().changeStates(states, 0, "", true, recommendationAdded) + private fun updateCodeWhisperer(sessionContext: SessionContext, states: InvocationContext, recommendationAdded: Boolean) { + CodeWhispererPopupManager.getInstance().changeStatesForShowing(sessionContext, states, recommendationAdded) } - private fun sendDiscardedUserDecisionEventForAll( - requestContext: RequestContext, - responseContext: ResponseContext, - recommendations: List, - ) { - val detailContexts = recommendations.map { - DetailContext("", it, it, true, false, "", getCompletionType(it)) - } - val recommendationContext = RecommendationContext(detailContexts, "", "", VisualPosition(0, 0)) + @RequiresEdt + private fun disposeJob(jobId: Int) { + ongoingRequests[jobId]?.let { Disposer.dispose(it) } + ongoingRequests.remove(jobId) + ongoingRequestsContext.remove(jobId) + } - CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll( - requestContext, - responseContext, - recommendationContext, - SessionContext(), - false - ) + @RequiresEdt + fun disposeDisplaySession(accept: Boolean) { + // avoid duplicate session disposal logic + if (sessionContext == null || sessionContext?.isDisposed() == true) return + + sessionContext?.let { + it.hasAccepted = accept + Disposer.dispose(it) + } + sessionContext = null + val jobIds = ongoingRequests.keys.toList() + jobIds.forEach { jobId -> disposeJob(jobId) } + ongoingRequests.clear() + ongoingRequestsContext.clear() } + fun getAllSuggestionsPreviewInfo() = + ongoingRequests.values.filterNotNull().flatMap { element -> + val context = element.recommendationContext + context.details.map { + PreviewContext(context.jobId, it, context.userInputSinceInvocation, context.typeahead) + } + } + + fun getAllPaginationSessions() = ongoingRequests + fun getRequestContext( triggerTypeInfo: TriggerTypeInfo, editor: Editor, project: Project, - psiFile: PsiFile, - latencyContext: LatencyContext, + psiFile: PsiFile ): RequestContext { // 1. file context val fileContext: FileContextInfo = runReadAction { FileContextProvider.getInstance(project).extractFileContext(editor, psiFile) } @@ -657,7 +704,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { // 5. customization val customizationArn = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)?.arn - return RequestContext(project, editor, triggerTypeInfo, caretPosition, fileContext, supplementalContext, connection, latencyContext, customizationArn) + return RequestContext(project, editor, triggerTypeInfo, caretPosition, fileContext, supplementalContext, connection, customizationArn) } fun validateResponse(response: GenerateCompletionsResponse): GenerateCompletionsResponse { @@ -683,26 +730,18 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { requestContext: RequestContext, responseContext: ResponseContext, recommendationContext: RecommendationContext, - popup: JBPopup, + coroutine: CoroutineScope ): InvocationContext { - addPopupChildDisposables(popup) // Creating a disposable for managing all listeners lifecycle attached to the popup. // previously(before pagination) we use popup as the parent disposable. // After pagination, listeners need to be updated as states are updated, for the same popup, // so disposable chain becomes popup -> disposable -> listeners updates, and disposable gets replaced on every // state update. - val states = InvocationContext(requestContext, responseContext, recommendationContext, popup) - Disposer.register(popup, states) - CodeWhispererPopupManager.getInstance().initPopupListener(states) - return states - } - - private fun addPopupChildDisposables(popup: JBPopup) { - codeInsightSettingsFacade.disableCodeInsightUntil(popup) - - Disposer.register(popup) { - CodeWhispererPopupManager.getInstance().reset() + val states = InvocationContext(requestContext, responseContext, recommendationContext) + Disposer.register(states) { + coroutine.cancel(CancellationException("Cancelling the current coroutine when the pagination session context is disposed")) } + return states } private fun logServiceInvocation( @@ -742,13 +781,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { return false } - if (CodeWhispererPopupManager.getInstance().hasConflictingPopups(editor)) { - LOG.debug { "Find other active popup windows before triggering CodeWhisperer, not invoking service" } - return false - } - - if (CodeWhispererInvocationStatus.getInstance().isPopupActive()) { - LOG.debug { "Find an existing CodeWhisperer popup window before triggering CodeWhisperer, not invoking service" } + if (CodeWhispererInvocationStatus.getInstance().isDisplaySessionActive()) { + LOG.debug { "Find an existing CodeWhisperer session before triggering CodeWhisperer, not invoking service" } return false } return true @@ -767,11 +801,13 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { companion object { private val LOG = getLogger() private const val MAX_REFRESH_ATTEMPT = 3 + private const val PAGINATION_REQUEST_COUNT_ALLOWED = 1 val CODEWHISPERER_CODE_COMPLETION_PERFORMED: Topic = Topic.create( "CodeWhisperer code completion service invoked", CodeWhispererCodeCompletionServiceListener::class.java ) + val DATA_KEY_SESSION = DataKey.create("codewhisperer.session") fun getInstance(): CodeWhispererService = service() const val KET_SESSION_ID = "x-amzn-SessionId" @@ -828,8 +864,7 @@ data class RequestContext( val fileContextInfo: FileContextInfo, private val supplementalContextDeferred: Deferred, val connection: ToolkitConnection?, - val latencyContext: LatencyContext, - val customizationArn: String?, + val customizationArn: String? ) { // TODO: should make the entire getRequestContext() suspend function instead of making supplemental context only var supplementalContext: SupplementalContextInfo? = null @@ -842,7 +877,6 @@ data class RequestContext( null } } - else -> field } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt index aefa1b51dad..ee72b7a321a 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt @@ -24,6 +24,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.customization.Code import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener @@ -86,7 +87,7 @@ abstract class CodeWhispererCodeCoverageTracker( conn.subscribe( CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED, object : CodeWhispererUserActionListener { - override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) { + override fun afterAccept(states: InvocationContext, previews: List, sessionContext: SessionContext, rangeMarker: RangeMarker) { if (states.requestContext.fileContextInfo.programmingLanguage != language) return rangeMarkers.add(rangeMarker) val originalRecommendation = extractRangeMarkerString(rangeMarker) ?: return diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceActionListener.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceActionListener.kt index 9576c002e05..1775631587b 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceActionListener.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceActionListener.kt @@ -5,14 +5,14 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow import com.intellij.openapi.editor.RangeMarker import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener class CodeWhispererCodeReferenceActionListener : CodeWhispererUserActionListener { - override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) { - val (project, editor) = states.requestContext - val manager = CodeWhispererCodeReferenceManager.getInstance(project) - manager.insertCodeReference(states, sessionContext.selectedIndex) - manager.addListeners(editor) + override fun afterAccept(states: InvocationContext, previews: List, sessionContext: SessionContext, rangeMarker: RangeMarker) { + val manager = CodeWhispererCodeReferenceManager.getInstance(sessionContext.project) + manager.insertCodeReference(states, previews, sessionContext.selectedIndex) + manager.addListeners(sessionContext.editor) } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceManager.kt index 7bb8dd16430..a3e2a90776b 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceManager.kt @@ -31,6 +31,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhisper import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.horizontalPanelConstraints import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext +import software.aws.toolkits.jetbrains.services.codewhisperer.model.PreviewContext import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.EDITOR_CODE_REFERENCE_HOVER import software.aws.toolkits.resources.message import javax.swing.JLabel @@ -109,11 +110,15 @@ class CodeWhispererCodeReferenceManager(private val project: Project) { } } - fun insertCodeReference(states: InvocationContext, selectedIndex: Int) { - val (requestContext, _, recommendationContext) = states - val (_, editor, _, caretPosition) = requestContext - val (_, detail, reformattedDetail) = recommendationContext.details[selectedIndex] - insertCodeReference(detail.content(), reformattedDetail.references(), editor, caretPosition, detail) + fun insertCodeReference(states: InvocationContext, previews: List, selectedIndex: Int) { + val detail = previews[selectedIndex].detail + insertCodeReference( + detail.recommendation.content(), + detail.reformatted.references(), + states.requestContext.editor, + states.requestContext.caretPosition, + detail.recommendation + ) } fun getReferenceLineNums(editor: Editor, start: Int, end: Int): String {