use project userData to store session instead of passing in local sta…

…tes to actions
aws · Sep 26, 2024 · 6370f4b · 6370f4b
1 parent 5b9903c
commit 6370f4b
Show file tree

Hide file tree

Showing 8 changed files with 14 additions and 32 deletions.
diff --git a/...ftware/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererAcceptAction.kt b/...ftware/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererAcceptAction.kt
@@ -9,14 +9,12 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.DumbAware
-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.CodeWhispererInvocationStatus
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.resources.message
 
 open class CodeWhispererAcceptAction(title: String = message("codewhisperer.inline.accept")) : AnAction(title), DumbAware {
-    var sessionContext: SessionContext? = null
-
     override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.EDT
 
     override fun update(e: AnActionEvent) {
@@ -25,10 +23,13 @@ open class CodeWhispererAcceptAction(title: String = message("codewhisperer.inli
     }
 
     override fun actionPerformed(e: AnActionEvent) {
-        val sessionContext = sessionContext ?: return
+        val sessionContext = e.project?.getUserData(CodeWhispererService.KEY_SESSION_CONTEXT) ?: return
         if (!CodeWhispererInvocationStatus.getInstance().isDisplaySessionActive()) return
         ApplicationManager.getApplication().messageBus.syncPublisher(
             CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED
         ).beforeAccept(sessionContext)
     }
 }
+
+// A same accept action but different key shortcut and different promoter logic
+class CodeWhispererForceAcceptAction(title: String = message("codewhisperer.inline.force.accept")) : CodeWhispererAcceptAction(title)
diff --git a/...e/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererForceAcceptAction.kt b/...e/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererForceAcceptAction.kt
diff --git a/.../aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererNavigateNextAction.kt b/.../aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererNavigateNextAction.kt
@@ -9,14 +9,12 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.DumbAware
-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.CodeWhispererInvocationStatus
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.resources.message
 
 class CodeWhispererNavigateNextAction : AnAction(message("codewhisperer.inline.navigate.next")), DumbAware {
-    var sessionContext: SessionContext? = null
-
     override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.EDT
 
     override fun update(e: AnActionEvent) {
@@ -25,7 +23,7 @@ class CodeWhispererNavigateNextAction : AnAction(message("codewhisperer.inline.n
     }
 
     override fun actionPerformed(e: AnActionEvent) {
-        val sessionContext = sessionContext ?: return
+        val sessionContext = e.project?.getUserData(CodeWhispererService.KEY_SESSION_CONTEXT) ?: return
         if (!CodeWhispererInvocationStatus.getInstance().isDisplaySessionActive()) return
         ApplicationManager.getApplication().messageBus.syncPublisher(
             CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED

diff --git a/.../aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererNavigatePrevAction.kt b/.../aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererNavigatePrevAction.kt
@@ -9,14 +9,12 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.DumbAware
-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.CodeWhispererInvocationStatus
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.resources.message
 
 class CodeWhispererNavigatePrevAction : AnAction(message("codewhisperer.inline.navigate.previous")), DumbAware {
-    var sessionContext: SessionContext? = null
-
     override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.EDT
 
     override fun update(e: AnActionEvent) {
@@ -25,7 +23,7 @@ class CodeWhispererNavigatePrevAction : AnAction(message("codewhisperer.inline.n
     }
 
     override fun actionPerformed(e: AnActionEvent) {
-        val sessionContext = sessionContext ?: return
+        val sessionContext = e.project?.getUserData(CodeWhispererService.KEY_SESSION_CONTEXT) ?: return
         if (!CodeWhispererInvocationStatus.getInstance().isDisplaySessionActive()) return
         ApplicationManager.getApplication().messageBus.syncPublisher(
             CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED

diff --git a/...software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt b/...software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt
@@ -33,6 +33,7 @@ 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.util.Key
 import com.intellij.openapi.wm.WindowManager
 import com.intellij.ui.ComponentUtil
 import com.intellij.ui.awt.RelativePoint
@@ -383,15 +384,7 @@ class CodeWhispererPopupManager {
     private fun setPopupActionHandlers(sessionContext: SessionContext) {
         val actionManager = EditorActionManager.getInstance()
 
-        // TODO: find a better way to pass in the local sessionContext for the handler to know the session state
-        val prevAction = ActionManager.getInstance().getAction("codewhisperer.inline.navigate.previous") as CodeWhispererNavigatePrevAction
-        prevAction.sessionContext = sessionContext
-        val nextAction = ActionManager.getInstance().getAction("codewhisperer.inline.navigate.next") as CodeWhispererNavigateNextAction
-        nextAction.sessionContext = sessionContext
-        val acceptAction = ActionManager.getInstance().getAction("codewhisperer.inline.accept") as CodeWhispererAcceptAction
-        acceptAction.sessionContext = sessionContext
-        val forceAcceptAction = ActionManager.getInstance().getAction("codewhisperer.inline.force.accept") as CodeWhispererForceAcceptAction
-        forceAcceptAction.sessionContext = sessionContext
+        sessionContext.project.putUserData(CodeWhispererService.KEY_SESSION_CONTEXT, sessionContext)
 
         setPopupTypedHandler(CodeWhispererPopupTypedHandler(TypedAction.getInstance().rawHandler, sessionContext), sessionContext)
         setPopupActionHandler(ACTION_EDITOR_ESCAPE, CodeWhispererPopupEscHandler(sessionContext), sessionContext)

diff --git a/...re/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererInvocationStatus.kt b/...re/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererInvocationStatus.kt
@@ -27,6 +27,7 @@ class CodeWhispererInvocationStatus {
 
     fun startInvocation() {
         isInvokingService.set(true)
+        ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(true)
         LOG.debug { "Starting CodeWhisperer invocation" }
     }
 

diff --git a/...s/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt b/...s/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt
@@ -17,7 +17,6 @@ import kotlin.math.min
 
 @Service
 class CodeWhispererRecommendationManager {
-    var states: MutableList<InvocationContext> = mutableListOf()
     fun reformatReference(requestContext: RequestContext, recommendation: Completion): Completion {
         // startOffset is the offset at the start of user input since invocation
         val invocationStartOffset = requestContext.caretPosition.offset

diff --git a/...rc/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/...rc/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt
@@ -17,6 +17,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.VisualPosition
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import com.intellij.util.concurrency.annotations.RequiresEdt
@@ -845,7 +846,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
             "CodeWhisperer intelliSense popup on hover",
             CodeWhispererIntelliSenseOnHoverListener::class.java
         )
-        val DATA_KEY_SESSION = DataKey.create<SessionContext>("codewhisperer.session")
+        val KEY_SESSION_CONTEXT = Key.create<SessionContext>("codewhisperer.session")
 
         fun getInstance(): CodeWhispererService = service()
         const val KET_SESSION_ID = "x-amzn-SessionId"