Track SessionContext instead of InvocationContext

1. Since the shift from displaying one trigger at a time
to displaying multiple triggers at a time, use SessionContext
as the disposable for the session states.
2. Introduce the concept of "display session" which can show
multiple triggers at a time and sessionContext is to hold states
for that session.
3. ongoingRequests will be a map of <jobId, InvocationContext>
where InvocationContext is the state of each trigger.
4. Managed these states globally by dynamically adding Trigger states
into ongoingRequests and whenever the user made a decision(accept or
reject), conclude decisions for all the ongoing triggers and clear
ongoingRequest, so that a new display session can start later.

Author: andrewyuq

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<CodeWhispererCustomization>
 
     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)

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)
             }
         }
     }

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]
 

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<CodeWhispererProgrammingLanguage>
     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<PreviewContext>, 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)

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<CodeWhispererImportAdderListener>()
-    override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) {
+    override fun afterAccept(states: InvocationContext, previews: List<PreviewContext>, 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)
     }
 }

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<Inlay<EditorCustomElementRenderer>>()
-    fun updateInlays(states: InvocationContext, chunks: List<RecommendationChunk>) {
-        val editor = states.requestContext.editor
+    fun updateInlays(sessionContext: SessionContext, chunks: List<RecommendationChunk>) {
         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)
             }
         }
     }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt

@@ -4,25 +4,37 @@
 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.popup.CodeWhispererPopupManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererIntelliSenseOnHoverListener
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 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.CodeWhispererUtil.setIntelliSensePopupAlpha
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CrossFileStrategy
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.SupplementalContextStrategy
 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 +93,19 @@ data class SupplementalContextInfo(
 }
 
 data class RecommendationContext(
-    val details: List<DetailContext>,
+    val details: MutableList<DetailContext>,
     val userInputOriginal: String,
     val userInputSinceInvocation: String,
-    val position: VisualPosition
+    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 +116,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<Int> = mutableSetOf(),
-    val isFirstTimeShowingPopup: Boolean = true,
+    var isFirstTimeShowingPopup: Boolean = true,
     var toBeRemovedHighlighter: RangeHighlighter? = null,
-    var insertEndOffset: Int = -1
-)
+    var insertEndOffset: Int = -1,
+    var popupDisplayOffset: 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 +175,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(

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)
     }
 }