From be2a6e297dde02a7fc5912046b8b54b8ed68f6bc Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Tue, 4 Feb 2025 15:12:03 -0800 Subject: [PATCH 1/4] delete dummy popup and set nullable popup --- .../inlay/CodeWhispererInlayManager.kt | 2 +- .../codewhisperer/model/CodeWhispererModel.kt | 2 +- .../popup/CodeWhispererPopupManager.kt | 26 +++++++++---------- .../popup/CodeWhispererUIChangeListener.kt | 6 +++-- .../handlers/CodeWhispererPopupEscHandler.kt | 2 +- .../service/CodeWhispererService.kt | 12 +++------ 6 files changed, 24 insertions(+), 26 deletions(-) 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..56c38688634 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 @@ -22,7 +22,7 @@ class CodeWhispererInlayManager { clearInlays() chunks.forEach { chunk -> - createCodeWhispererInlays(editor, chunk.inlayOffset, chunk.text, states.popup) + states.popup?.let { createCodeWhispererInlays(editor, chunk.inlayOffset, chunk.text, 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 af2bd28c33e..b7cda08c3d9 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 @@ -216,7 +216,7 @@ data class InvocationContext( val requestContext: RequestContext, val responseContext: ResponseContext, val recommendationContext: RecommendationContext, - val popup: JBPopup, + val popup: JBPopup? = null, ) : Disposable { override fun dispose() {} } 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 0c38e1eaeb6..55be5f01e1c 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 @@ -141,7 +141,7 @@ class CodeWhispererPopupManager { val startOffset = states.requestContext.caretPosition.offset val currOffset = states.requestContext.editor.caretModel.offset if (startOffset > currOffset) { - cancelPopup(popup) + popup?.let { cancelPopup(popup) } return } @@ -149,7 +149,7 @@ class CodeWhispererPopupManager { val prefix = states.requestContext.editor.document.charsSequence .substring(startOffset, currOffset) if (prefix.length < userInputOriginal.length) { - cancelPopup(popup) + popup?.let { cancelPopup(popup) } return } else { prefix.substring(userInputOriginal.length) @@ -172,7 +172,7 @@ class CodeWhispererPopupManager { ) if (selectedIndex == -1 || !isValidRecommendation(details[selectedIndex], userInput, typeaheadOriginal)) { LOG.debug { "None of the recommendation is valid at this point, cancelling the popup" } - cancelPopup(popup) + popup?.let { cancelPopup(popup) } return } val typeahead = resolveTypeahead(states, selectedIndex, typeaheadOriginal) @@ -245,7 +245,7 @@ class CodeWhispererPopupManager { states.requestContext.latencyContext.getPerceivedLatency(states.requestContext.triggerTypeInfo.triggerType) } if (!isRecommendationAdded) { - showPopup(states, sessionContext, states.popup, visible = sessionContext.isPopupShowing) + states.popup?.let { showPopup(states, sessionContext, it, visible = sessionContext.isPopupShowing) } } } @@ -423,8 +423,8 @@ class CodeWhispererPopupManager { private fun addPopupListener(states: InvocationContext) { val listener = CodeWhispererPopupListener(states) - states.popup.addListener(listener) - Disposer.register(states) { states.popup.removeListener(listener) } + states.popup?.addListener(listener) + Disposer.register(states) { states.popup?.removeListener(listener) } } private fun addMessageSubscribers(states: InvocationContext) { @@ -460,7 +460,7 @@ class CodeWhispererPopupManager { dontClosePopupAndRun { CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext) } - closePopup(states.popup) + states.popup?.let { closePopup(it) } if (sessionContext.selectedIndex == 0) { CodeWhispererService.getInstance().promoteNextInvocationIfAvailable() } @@ -512,7 +512,7 @@ class CodeWhispererPopupManager { val codewhispererSelectionListener: SelectionListener = object : SelectionListener { override fun selectionChanged(event: SelectionEvent) { if (allowEditsDuringSuggestionPreview.availablePermits == MAX_EDIT_SOURCE_DURING_SUGGESTION_PREVIEW) { - cancelPopup(states.popup) + states.popup?.let { cancelPopup(it) } } super.selectionChanged(event) } @@ -528,7 +528,7 @@ class CodeWhispererPopupManager { if (editor.caretModel.offset == event.offset) { changeStates(states, 0) } else if (allowEditsDuringSuggestionPreview.availablePermits == MAX_EDIT_SOURCE_DURING_SUGGESTION_PREVIEW) { - cancelPopup(states.popup) + states.popup?.let { cancelPopup(it) } } } } @@ -537,7 +537,7 @@ class CodeWhispererPopupManager { val codewhispererCaretListener: CaretListener = object : CaretListener { override fun caretPositionChanged(event: CaretEvent) { if (allowEditsDuringSuggestionPreview.availablePermits == MAX_EDIT_SOURCE_DURING_SUGGESTION_PREVIEW) { - cancelPopup(states.popup) + states.popup?.let { cancelPopup(it) } } super.caretPositionChanged(event) } @@ -550,11 +550,11 @@ class CodeWhispererPopupManager { val window = ComponentUtil.getWindow(editorComponent) val windowListener: ComponentListener = object : ComponentAdapter() { override fun componentMoved(event: ComponentEvent) { - cancelPopup(states.popup) + states.popup?.let { cancelPopup(it) } } override fun componentShown(e: ComponentEvent?) { - cancelPopup(states.popup) + states.popup?.let { cancelPopup(it) } super.componentShown(e) } } @@ -565,7 +565,7 @@ class CodeWhispererPopupManager { val suggestionHoverEnterListener: EditorMouseMotionListener = object : EditorMouseMotionListener { override fun mouseMoved(e: EditorMouseEvent) { if (e.inlay != null) { - showPopup(states, sessionContext, states.popup, visible = true) + states.popup?.let { showPopup(states, sessionContext, it, visible = true) } } else { bringSuggestionInlayToFront(editor, states.popup, sessionContext, opposite = true) } 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 6de006145a6..302736a47d1 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 @@ -61,8 +61,10 @@ class CodeWhispererUIChangeListener : CodeWhispererPopupStateChangeListener { }, HighlighterTargetArea.EXACT_RANGE ) - Disposer.register(states.popup) { - editor.markupModel.removeHighlighter(rangeHighlighter) + states.popup?.let { + Disposer.register(it) { + editor.markupModel.removeHighlighter(rangeHighlighter) + } } sessionContext.toBeRemovedHighlighter = rangeHighlighter } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/handlers/CodeWhispererPopupEscHandler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/handlers/CodeWhispererPopupEscHandler.kt index e4dbfd32e9f..dc63b0b2b2f 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/handlers/CodeWhispererPopupEscHandler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/handlers/CodeWhispererPopupEscHandler.kt @@ -13,7 +13,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe class CodeWhispererPopupEscHandler(states: InvocationContext) : CodeWhispererEditorActionHandler(states) { override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext?) { - CodeWhispererPopupManager.getInstance().cancelPopup(states.popup) + states.popup?.let { CodeWhispererPopupManager.getInstance().cancelPopup(it) } } } 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 1b370cb0dbb..88e57ebb15b 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 @@ -16,7 +16,6 @@ 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.ui.popup.JBPopupFactory import com.intellij.openapi.util.Disposer import com.intellij.psi.PsiDocumentManager import com.intellij.psi.PsiFile @@ -299,12 +298,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { } } val nextRecommendationContext = RecommendationContext(detailContexts, "", "", newVisualPosition) - val newPopup = withContext(EDT) { - JBPopupFactory.getInstance().createMessage("Dummy popup") - } // send userDecision and trigger decision when next recommendation haven't been seen - if (currStates.popup.isDisposed) { + if (currStates.popup?.isDisposed == true) { CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll( nextRequestContext, nextResponseContext, @@ -313,7 +309,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { false ) } else { - nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, newPopup) + nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, null) } LOG.debug { "Prefetched next invocation stored in nextInvocationContext" } } @@ -709,7 +705,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { val updatedStates = states.copy( recommendationContext = recommendationContext.copy(details = details + newDetailContexts) ) - Disposer.register(states.popup, updatedStates) + states.popup?.let { Disposer.register(it, updatedStates) } CodeWhispererPopupManager.getInstance().initPopupListener(updatedStates) return updatedStates } @@ -744,7 +740,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { cs.launch { val newPopup = CodeWhispererPopupManager.getInstance().initPopup() val updatedNextStates = nextStates.copy(popup = newPopup).also { - addPopupChildDisposables(it.requestContext.project, it.requestContext.editor, it.popup) + it.popup?.let { it1 -> addPopupChildDisposables(it.requestContext.project, it.requestContext.editor, it1) } Disposer.register(newPopup, it) } CodeWhispererPopupManager.getInstance().initPopupListener(updatedNextStates) From ab5ef788248720f084e05ca1d0b8fcfb25a90a8e Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Mon, 10 Feb 2025 13:40:25 -0800 Subject: [PATCH 2/4] check session id --- .../services/codewhisperer/service/CodeWhispererService.kt | 6 ++++++ 1 file changed, 6 insertions(+) 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 88e57ebb15b..547d2a88207 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 @@ -26,6 +26,7 @@ import kotlinx.coroutines.Deferred import kotlinx.coroutines.Job import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.async +import kotlinx.coroutines.cancel import kotlinx.coroutines.delay import kotlinx.coroutines.isActive import kotlinx.coroutines.launch @@ -267,6 +268,11 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble() val requestId = nextResponse.responseMetadata().requestId() val sessionId = nextResponse.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0] + if (sessionId != currStates.responseContext.sessionId) { + LOG.debug { "Session id mismatch, expected: ${currStates.responseContext.sessionId}, actual: $sessionId" } + this.cancel("Session mismatch") + return@launch + } nextRequestContext.latencyContext.apply { codewhispererPostprocessingStart = System.nanoTime() From 83e38739d690d2a98fc9eb34a686ec52b22884cd Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Mon, 10 Feb 2025 13:58:27 -0800 Subject: [PATCH 3/4] wrap with ? call for nullable variable --- .../services/codewhisperer/popup/CodeWhispererPopupManager.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 1e1e47e9cf2..da981137b02 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 @@ -243,7 +243,7 @@ class CodeWhispererPopupManager { states.requestContext.latencyContext.getPerceivedLatency(states.requestContext.triggerTypeInfo.triggerType) } if (!isRecommendationAdded) { - showPopup(states, sessionContext, states.popup, visible = sessionContext.isPopupShowing) + states.popup?.let { showPopup(states, sessionContext, it, visible = sessionContext.isPopupShowing) } } } From 82c00a92db1e068191254269fea7d5decec88563 Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Mon, 10 Feb 2025 15:43:25 -0800 Subject: [PATCH 4/4] add changlog --- .../feature-5cdc6e55-3a3f-4a31-8499-543cf249c734.json | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 .changes/next-release/feature-5cdc6e55-3a3f-4a31-8499-543cf249c734.json diff --git a/.changes/next-release/feature-5cdc6e55-3a3f-4a31-8499-543cf249c734.json b/.changes/next-release/feature-5cdc6e55-3a3f-4a31-8499-543cf249c734.json new file mode 100644 index 00000000000..d93d942e593 --- /dev/null +++ b/.changes/next-release/feature-5cdc6e55-3a3f-4a31-8499-543cf249c734.json @@ -0,0 +1,4 @@ +{ + "type" : "feature", + "description" : "Inline suggestions: Enable Pre-fetch recommendations to reduce suggestion latency for complex scenarios" +} \ No newline at end of file