fix: should keep reporting UTDE telemetry if there are still pending Edits suggestions (aws#2051)

Author: atonaamz

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts

@@ -681,6 +681,7 @@ export const CodewhispererServerFactory =
                     partialResultToken: suggestionResponse.responseContext.nextToken,
                 }
             } else {
+                session.hasEditsPending = suggestionResponse.responseContext.nextToken ? true : false
                 return {
                     items: suggestionResponse.suggestions
                         .map(suggestion => {
@@ -867,7 +868,13 @@ export const CodewhispererServerFactory =
             if (firstCompletionDisplayLatency) emitPerceivedLatencyTelemetry(telemetry, session)
 
             // Always emit user trigger decision at session close
-            sessionManager.closeSession(session)
+            // Close session unless Edit suggestion was accepted with more pending
+            const shouldKeepSessionOpen =
+                session.suggestionType === SuggestionType.EDIT && isAccepted && session.hasEditsPending
+
+            if (!shouldKeepSessionOpen) {
+                sessionManager.closeSession(session)
+            }
             const streakLength = editsEnabled ? sessionManager.getAndUpdateStreakLength(isAccepted) : 0
             await emitUserTriggerDecisionTelemetry(
                 telemetry,

server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts

@@ -6,7 +6,12 @@ import {
 } from '@aws/language-server-runtimes/server-interface'
 import { v4 as uuidv4 } from 'uuid'
 import { CodewhispererAutomatedTriggerType, CodewhispererTriggerType } from '../auto-trigger/autoTrigger'
-import { GenerateSuggestionsRequest, ResponseContext, Suggestion } from '../../../shared/codeWhispererService'
+import {
+    GenerateSuggestionsRequest,
+    ResponseContext,
+    Suggestion,
+    SuggestionType,
+} from '../../../shared/codeWhispererService'
 import { CodewhispererLanguage } from '../../../shared/languageDetection'
 import { CodeWhispererSupplementalContext } from '../../../shared/models/model'
 import { Logging } from '@aws/language-server-runtimes/server-interface'
@@ -75,6 +80,9 @@ export class CodeWhispererSession {
     includeImportsWithSuggestions?: boolean
     codewhispererSuggestionImportCount: number = 0
     suggestionType?: string
+    hasEditsPending?: boolean = false
+    // Track the most recent itemId for paginated Edit suggestions
+    recentItemId?: string
 
     constructor(data: SessionData) {
         this.id = this.generateSessionId()
@@ -153,7 +161,11 @@ export class CodeWhispererSession {
         typeaheadLength?: number
     ) {
         // Skip if session results were already recorded for session of session is closed
-        if (this.state === 'CLOSED' || this.state === 'DISCARD' || this.completionSessionResult) {
+        if (
+            this.state === 'CLOSED' ||
+            this.state === 'DISCARD' ||
+            (this.completionSessionResult && this.suggestionType === SuggestionType.COMPLETION)
+        ) {
             return
         }
 
@@ -195,6 +207,8 @@ export class CodeWhispererSession {
                 // No recommendation was accepted, but user have seen this suggestion
                 this.setSuggestionState(itemId, 'Reject')
             }
+
+            this.recentItemId = itemId
         }
 
         this.firstCompletionDisplayLatency = firstCompletionDisplayLatency
@@ -240,6 +254,24 @@ export class CodeWhispererSession {
         }
         return isEmpty ? 'Empty' : 'Discard'
     }
+
+    /**
+     * Determines trigger decision based on the most recent user action.
+     * Uses the last processed itemId to determine the overall session decision.
+     */
+    getLatestUserTriggerDecision(): UserTriggerDecision | undefined {
+        // Force Discard trigger decision when session was explicitly discarded by server
+        if (this.state === 'DISCARD') {
+            return 'Discard'
+        }
+
+        if (!this.recentItemId) return
+
+        const state = this.getSuggestionState(this.recentItemId)
+        if (state === 'Accept') return 'Accept'
+        if (state === 'Reject') return 'Reject'
+        return state === 'Empty' ? 'Empty' : 'Discard'
+    }
 }
 
 export class SessionManager {

server/aws-lsp-codewhisperer/src/language-server/inline-completion/telemetry.ts

@@ -4,6 +4,7 @@ import { CodeWhispererSession } from './session/sessionManager'
 import { CodeWhispererPerceivedLatencyEvent, CodeWhispererServiceInvocationEvent } from '../../shared/telemetry/types'
 import { getCompletionType, isAwsError } from '../../shared/utils'
 import { TelemetryService } from '../../shared/telemetry/telemetryService'
+import { SuggestionType } from '../../shared/codeWhispererService'
 
 export const emitServiceInvocationTelemetry = (
     telemetry: Telemetry,
@@ -121,8 +122,15 @@ export const emitUserTriggerDecisionTelemetry = async (
         return
     }
 
+    // Edits show one suggestion sequentially (with pagination), so use latest itemId state;
+    // Completions show multiple suggestions together, so aggregate all states
+    const userTriggerDecision =
+        session.suggestionType === SuggestionType.EDIT
+            ? session.getLatestUserTriggerDecision()
+            : session.getAggregatedUserTriggerDecision()
+
     // Can not emit previous trigger decision if it's not available on the session
-    if (!session.getAggregatedUserTriggerDecision()) {
+    if (!userTriggerDecision) {
         return
     }
 
@@ -137,7 +145,13 @@ export const emitUserTriggerDecisionTelemetry = async (
         streakLength
     )
 
-    session.reportedUserDecision = true
+    // Mark telemetry as complete unless Edit suggestion was accepted with more pending
+    const hasPendingEditTelemetry =
+        session.suggestionType === SuggestionType.EDIT && session.acceptedSuggestionId && session.hasEditsPending
+
+    if (!hasPendingEditTelemetry) {
+        session.reportedUserDecision = true
+    }
 }
 
 export const emitAggregatedUserTriggerDecisionTelemetry = (

server/aws-lsp-codewhisperer/src/shared/telemetry/telemetryService.ts

@@ -117,7 +117,13 @@ export class TelemetryService {
 
     private getSuggestionState(session: CodeWhispererSession): SuggestionState {
         let suggestionState: SuggestionState
-        switch (session.getAggregatedUserTriggerDecision()) {
+        // Edits show one suggestion sequentially (with pagination), so use latest itemId state;
+        // Completions show multiple suggestions together, so aggregate all states
+        const userTriggerDecision =
+            session.suggestionType === SuggestionType.EDIT
+                ? session.getLatestUserTriggerDecision()
+                : session.getAggregatedUserTriggerDecision()
+        switch (userTriggerDecision) {
             case 'Accept':
                 suggestionState = 'ACCEPT'
                 break