Skip to content

Track SessionContext instead of InvocationContext #4916

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 4, 2024
Merged

Track SessionContext instead of InvocationContext #4916

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2a414b3
Track SessionContext instead of InvocationContext
andrewyuq Sep 25, 2024
d37fec2
address comments
andrewyuq Sep 26, 2024
156159e
Merge remote-tracking branch 'aws/auto-trigger-revamp' into cw-teleme…
andrewyuq Oct 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions ...e/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -92,6 +93,7 @@ interface CodeWhispererClientAdaptor : Disposable {
fun listAvailableCustomizations(): List<CodeWhispererCustomization>

fun sendUserTriggerDecisionTelemetry(
sessionContext: SessionContext,
requestContext: RequestContext,
responseContext: ResponseContext,
completionType: CodewhispererCompletionType,
Expand Down Expand Up @@ -293,6 +295,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
}

override fun sendUserTriggerDecisionTelemetry(
sessionContext: SessionContext,
requestContext: RequestContext,
responseContext: ResponseContext,
completionType: CodewhispererCompletionType,
Expand All @@ -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)
Expand Down
24 changes: 15 additions & 9 deletions ...ftware/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -67,7 +73,7 @@ class CodeWhispererEditorManager {

ApplicationManager.getApplication().messageBus.syncPublisher(
CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED,
).afterAccept(states, sessionContext, rangeMarker)
).afterAccept(states, previews, sessionContext, rangeMarker)
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion .../software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down
6 changes: 3 additions & 3 deletions ...are/aws/toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdder.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 3 additions & 2 deletions ...toolkits/jetbrains/services/codewhisperer/importadder/CodeWhispererImportAdderListener.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
}
}
16 changes: 7 additions & 9 deletions ...software/aws/toolkits/jetbrains/services/codewhisperer/inlay/CodeWhispererInlayManager.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
}
}

Expand All @@ -73,7 +71,7 @@ class CodeWhispererInlayManager {
)
blockInlay?.let {
existingInlays.add(it)
Disposer.register(popup, it)
Disposer.register(sessionContext, it)
}
}
}
Expand Down
70 changes: 61 additions & 9 deletions ...ty/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,25 +4,33 @@
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
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(
Expand Down Expand Up @@ -81,10 +89,19 @@ data class SupplementalContextInfo(
}

data class RecommendationContext(
val details: List<DetailContext>,
val details: MutableList<DetailContext>,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not seeing why need mutable

Copy link
Contributor Author

@andrewyuq andrewyuq Sep 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now each trigger will only have one state InvocationContext and is stored in ongoingRequests, when subsequent suggestions sent back to client, they will be added to this mutable details, and other parts of InvocationContext remains the same.

private fun updateStates(
        states: InvocationContext,
        response: GenerateCompletionsResponse
    ): InvocationContext {
        val recommendationContext = states.recommendationContext
        val newDetailContexts = CodeWhispererRecommendationManager.getInstance().buildDetailContext(
            states.requestContext,
            recommendationContext.userInputSinceInvocation,
            response.completions(),
            response.responseMetadata().requestId()
        )

        recommendationContext.details.addAll(newDetailContexts)
        return states
    }

otherwise we will have to replace InvocationContext which I feel is a bit unnecessary.

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,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dumb q, what's the difference between these 2 typeaheads

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

previously we have typeahead and typeaheadoriginal for the purpose of accounting leading spaces in the typeahead,
for example, typeaheadoriginal will be System and typeahead will be System, and then we use typeahead to do the suggestion matching logic. This is to preserve the suggestions as much as possible.

Now we trigger more so we have more display oppotunities, I feel like this becomes a bit unnecessary and makes the logic complicated, so remove typeahead original to make it no longer accomodate leading spaces.

)

data class DetailContext(
Expand All @@ -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<Int> = 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
Comment on lines +143 to +149
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

popup manager and sessioncontext are so tied together that it feels like it should just be handled together

Copy link
Contributor Author

@andrewyuq andrewyuq Sep 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

popup now becomes a state of a session (1 session 1 popup) so I put it as a field of SessionContext, and then I let disposing SessionContext handles everything so I don't have to call

CWService.getInstance().disposeDisplaySession()
popup.cancel()/close()

separately.

CodeWhispererInvocationStatus.getInstance().finishInvocation()
isDisposed = true
}

fun isDisposed() = isDisposed
}

data class RecommendationChunk(
val text: String,
Expand All @@ -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(
Expand Down
25 changes: 3 additions & 22 deletions ...oftware/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
Loading