fix: adding streakTracker to track streakLength across Completions and Edits (aws#2147)

Author: atonaamz

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

@@ -46,6 +46,7 @@ import { UserWrittenCodeTracker } from '../../shared/userWrittenCodeTracker'
 import { RecentEditTracker, RecentEditTrackerDefaultConfig } from './tracker/codeEditTracker'
 import { CursorTracker } from './tracker/cursorTracker'
 import { RejectedEditTracker, DEFAULT_REJECTED_EDIT_TRACKER_CONFIG } from './tracker/rejectedEditTracker'
+import { StreakTracker } from './tracker/streakTracker'
 import { getAddedAndDeletedLines, getCharacterDifferences } from './diffUtils'
 import {
     emitPerceivedLatencyTelemetry,
@@ -129,6 +130,7 @@ export const CodewhispererServerFactory =
         const recentEditTracker = RecentEditTracker.getInstance(logging, RecentEditTrackerDefaultConfig)
         const cursorTracker = CursorTracker.getInstance()
         const rejectedEditTracker = RejectedEditTracker.getInstance(logging, DEFAULT_REJECTED_EDIT_TRACKER_CONFIG)
+        const streakTracker = StreakTracker.getInstance()
         let editsEnabled = false
         let isOnInlineCompletionHandlerInProgress = false
 
@@ -317,9 +319,7 @@ export const CodewhispererServerFactory =
                         // for the previous trigger
                         if (ideCategory !== 'JETBRAINS') {
                             completionSessionManager.discardSession(currentSession)
-                            const streakLength = editsEnabled
-                                ? completionSessionManager.getAndUpdateStreakLength(false)
-                                : 0
+                            const streakLength = editsEnabled ? streakTracker.getAndUpdateStreakLength(false) : 0
                             await emitUserTriggerDecisionTelemetry(
                                 telemetry,
                                 telemetryService,
@@ -405,7 +405,7 @@ export const CodewhispererServerFactory =
             if (session.discardInflightSessionOnNewInvocation) {
                 session.discardInflightSessionOnNewInvocation = false
                 completionSessionManager.discardSession(session)
-                const streakLength = editsEnabled ? completionSessionManager.getAndUpdateStreakLength(false) : 0
+                const streakLength = editsEnabled ? streakTracker.getAndUpdateStreakLength(false) : 0
                 await emitUserTriggerDecisionTelemetry(
                     telemetry,
                     telemetryService,
@@ -664,7 +664,7 @@ export const CodewhispererServerFactory =
 
             // Always emit user trigger decision at session close
             sessionManager.closeSession(session)
-            const streakLength = editsEnabled ? sessionManager.getAndUpdateStreakLength(isAccepted) : 0
+            const streakLength = editsEnabled ? streakTracker.getAndUpdateStreakLength(isAccepted) : 0
             await emitUserTriggerDecisionTelemetry(
                 telemetry,
                 telemetryService,

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

@@ -36,12 +36,14 @@ import { getErrorMessage, hasConnectionExpired } from '../../shared/utils'
 import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../shared/amazonQServiceManager/errors'
 import { DocumentChangedListener } from './documentChangedListener'
 import { EMPTY_RESULT, EDIT_DEBOUNCE_INTERVAL_MS } from './constants'
+import { StreakTracker } from './tracker/streakTracker'
 
 export class EditCompletionHandler {
     private readonly editsEnabled: boolean
     private debounceTimeout: NodeJS.Timeout | undefined
     private isWaiting: boolean = false
     private hasDocumentChangedSinceInvocation: boolean = false
+    private readonly streakTracker: StreakTracker
 
     constructor(
         readonly logging: Logging,
@@ -60,6 +62,7 @@ export class EditCompletionHandler {
         this.editsEnabled =
             this.clientMetadata.initializationOptions?.aws?.awsClientCapabilities?.textDocument
                 ?.inlineCompletionWithReferences?.inlineEditSupport ?? false
+        this.streakTracker = StreakTracker.getInstance()
     }
 
     get codeWhispererService() {
@@ -264,7 +267,7 @@ export class EditCompletionHandler {
         if (currentSession && currentSession.state === 'ACTIVE') {
             // Emit user trigger decision at session close time for active session
             this.sessionManager.discardSession(currentSession)
-            const streakLength = this.editsEnabled ? this.sessionManager.getAndUpdateStreakLength(false) : 0
+            const streakLength = this.editsEnabled ? this.streakTracker.getAndUpdateStreakLength(false) : 0
             await emitUserTriggerDecisionTelemetry(
                 this.telemetry,
                 this.telemetryService,
@@ -335,7 +338,7 @@ export class EditCompletionHandler {
         if (session.discardInflightSessionOnNewInvocation) {
             session.discardInflightSessionOnNewInvocation = false
             this.sessionManager.discardSession(session)
-            const streakLength = this.editsEnabled ? this.sessionManager.getAndUpdateStreakLength(false) : 0
+            const streakLength = this.editsEnabled ? this.streakTracker.getAndUpdateStreakLength(false) : 0
             await emitUserTriggerDecisionTelemetry(
                 this.telemetry,
                 this.telemetryService,
@@ -359,7 +362,7 @@ export class EditCompletionHandler {
                 this.telemetryService,
                 session,
                 this.documentChangedListener.timeSinceLastUserModification,
-                this.editsEnabled ? this.sessionManager.getAndUpdateStreakLength(false) : 0
+                this.editsEnabled ? this.streakTracker.getAndUpdateStreakLength(false) : 0
             )
             return EMPTY_RESULT
         }

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

@@ -673,46 +673,4 @@ describe('SessionManager', function () {
             assert.equal(session.getSuggestionState('id4'), 'Discard')
         })
     })
-
-    describe('getAndUpdateStreakLength()', function () {
-        it('should return 0 if user rejects suggestion A', function () {
-            const manager = SessionManager.getInstance()
-
-            assert.equal(manager.getAndUpdateStreakLength(false), -1)
-            assert.equal(manager.streakLength, 0)
-        })
-
-        it('should return -1 for A and 1 for B if user accepts suggestion A and rejects B', function () {
-            const manager = SessionManager.getInstance()
-
-            assert.equal(manager.getAndUpdateStreakLength(true), -1)
-            assert.equal(manager.streakLength, 1)
-            assert.equal(manager.getAndUpdateStreakLength(false), 1)
-            assert.equal(manager.streakLength, 0)
-        })
-
-        it('should return -1 for A, -1 for B, and 2 for C if user accepts A, accepts B, and rejects C', function () {
-            const manager = SessionManager.getInstance()
-
-            assert.equal(manager.getAndUpdateStreakLength(true), -1)
-            assert.equal(manager.streakLength, 1)
-            assert.equal(manager.getAndUpdateStreakLength(true), -1)
-            assert.equal(manager.streakLength, 2)
-            assert.equal(manager.getAndUpdateStreakLength(false), 2)
-            assert.equal(manager.streakLength, 0)
-        })
-
-        it('should return -1 for A, -1 for B, and 1 for C if user accepts A, make an edit, accepts B, and rejects C', function () {
-            const manager = SessionManager.getInstance()
-
-            assert.equal(manager.getAndUpdateStreakLength(true), -1)
-            assert.equal(manager.streakLength, 1)
-            assert.equal(manager.getAndUpdateStreakLength(false), 1)
-            assert.equal(manager.streakLength, 0)
-            assert.equal(manager.getAndUpdateStreakLength(true), -1)
-            assert.equal(manager.streakLength, 1)
-            assert.equal(manager.getAndUpdateStreakLength(false), 1)
-            assert.equal(manager.streakLength, 0)
-        })
-    })
 })

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

@@ -282,7 +282,6 @@ export class SessionManager {
     private currentSession?: CodeWhispererSession
     private sessionsLog: CodeWhispererSession[] = []
     private maxHistorySize = 5
-    streakLength: number = 0
     // TODO, for user decision telemetry: accepted suggestions (not necessarily the full corresponding session) should be stored for 5 minutes
 
     private constructor() {}
@@ -362,16 +361,4 @@ export class SessionManager {
             this.currentSession.activate()
         }
     }
-
-    getAndUpdateStreakLength(isAccepted: boolean | undefined): number {
-        if (!isAccepted && this.streakLength != 0) {
-            const currentStreakLength = this.streakLength
-            this.streakLength = 0
-            return currentStreakLength
-        } else if (isAccepted) {
-            // increment streakLength everytime a suggestion is accepted.
-            this.streakLength = this.streakLength + 1
-        }
-        return -1
-    }
 }

server/aws-lsp-codewhisperer/src/language-server/inline-completion/tracker/streakTracker.test.ts

@@ -0,0 +1,85 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import { StreakTracker } from './streakTracker'
+
+describe('StreakTracker', function () {
+    let tracker: StreakTracker
+
+    beforeEach(function () {
+        StreakTracker.reset()
+        tracker = StreakTracker.getInstance()
+    })
+
+    afterEach(function () {
+        StreakTracker.reset()
+    })
+
+    describe('getInstance', function () {
+        it('should return the same instance (singleton)', function () {
+            const instance1 = StreakTracker.getInstance()
+            const instance2 = StreakTracker.getInstance()
+            assert.strictEqual(instance1, instance2)
+        })
+
+        it('should create new instance after reset', function () {
+            const instance1 = StreakTracker.getInstance()
+            StreakTracker.reset()
+            const instance2 = StreakTracker.getInstance()
+            assert.notStrictEqual(instance1, instance2)
+        })
+    })
+
+    describe('getAndUpdateStreakLength', function () {
+        it('should return -1 for undefined input', function () {
+            const result = tracker.getAndUpdateStreakLength(undefined)
+            assert.strictEqual(result, -1)
+        })
+
+        it('should return -1 and increment streak on acceptance', function () {
+            const result = tracker.getAndUpdateStreakLength(true)
+            assert.strictEqual(result, -1)
+        })
+
+        it('should return -1 for rejection with zero streak', function () {
+            const result = tracker.getAndUpdateStreakLength(false)
+            assert.strictEqual(result, -1)
+        })
+
+        it('should return previous streak on rejection after acceptances', function () {
+            tracker.getAndUpdateStreakLength(true)
+            tracker.getAndUpdateStreakLength(true)
+            tracker.getAndUpdateStreakLength(true)
+
+            const result = tracker.getAndUpdateStreakLength(false)
+            assert.strictEqual(result, 3)
+        })
+
+        it('should handle acceptance after rejection', function () {
+            tracker.getAndUpdateStreakLength(true)
+            tracker.getAndUpdateStreakLength(true)
+
+            const resetResult = tracker.getAndUpdateStreakLength(false)
+            assert.strictEqual(resetResult, 2)
+
+            tracker.getAndUpdateStreakLength(true)
+            const newResult = tracker.getAndUpdateStreakLength(true)
+            assert.strictEqual(newResult, -1)
+        })
+    })
+
+    describe('cross-instance consistency', function () {
+        it('should maintain state across getInstance calls', function () {
+            const tracker1 = StreakTracker.getInstance()
+            tracker1.getAndUpdateStreakLength(true)
+            tracker1.getAndUpdateStreakLength(true)
+
+            const tracker2 = StreakTracker.getInstance()
+            const result = tracker2.getAndUpdateStreakLength(false)
+            assert.strictEqual(result, 2)
+        })
+    })
+})

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

@@ -0,0 +1,42 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Tracks acceptance streak across both completion and edit suggestion types.
+ * Shared singleton to maintain consistent streak count between different code paths.
+ */
+export class StreakTracker {
+    private static _instance?: StreakTracker
+    private streakLength: number = 0
+
+    private constructor() {}
+
+    public static getInstance(): StreakTracker {
+        if (!StreakTracker._instance) {
+            StreakTracker._instance = new StreakTracker()
+        }
+        return StreakTracker._instance
+    }
+
+    public static reset() {
+        StreakTracker._instance = undefined
+    }
+
+    /**
+     * Updates and returns the current streak length based on acceptance status.
+     * @param isAccepted Whether the suggestion was accepted
+     * @returns Current streak length before update, or -1 if no change
+     */
+    public getAndUpdateStreakLength(isAccepted: boolean | undefined): number {
+        if (!isAccepted && this.streakLength !== 0) {
+            const currentStreakLength = this.streakLength
+            this.streakLength = 0
+            return currentStreakLength
+        } else if (isAccepted) {
+            this.streakLength += 1
+        }
+        return -1
+    }
+}