fix(amazonq): address pr feedback

Author: zixlin7

packages/amazonq/src/app/inline/completion.ts

@@ -16,10 +16,7 @@ import {
     Disposable,
 } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
-import {
-    logInlineCompletionSessionResultsNotificationType,
-    LogInlineCompletionSessionResultsParams,
-} from '@aws/language-server-runtimes/protocol'
+import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-runtimes/protocol'
 import { SessionManager } from './sessionManager'
 import { RecommendationService } from './recommendationService'
 
@@ -56,10 +53,19 @@ export class InlineCompletionManager implements Disposable {
     private disposable: Disposable
     private inlineCompletionProvider: AmazonQInlineCompletionItemProvider
     private languageClient: LanguageClient
+    private sessionManager: SessionManager
+    private recommendationService: RecommendationService
+    private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults'
 
     constructor(languageClient: LanguageClient) {
         this.languageClient = languageClient
-        this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(languageClient)
+        this.sessionManager = new SessionManager()
+        this.recommendationService = new RecommendationService(this.sessionManager)
+        this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(
+            languageClient,
+            this.recommendationService,
+            this.sessionManager
+        )
         this.disposable = languages.registerInlineCompletionItemProvider(
             CodewhispererInlineCompletionLanguages,
             this.inlineCompletionProvider
@@ -92,16 +98,16 @@ export class InlineCompletionManager implements Disposable {
                 totalSessionDisplayTime: Date.now() - requestStartTime,
                 firstCompletionDisplayLatency: firstCompletionDisplayLatency,
             }
-            this.languageClient.sendNotification(logInlineCompletionSessionResultsNotificationType as any, params)
+            this.languageClient.sendNotification(this.logSessionResultMessageName, params)
             this.disposable.dispose()
             this.disposable = languages.registerInlineCompletionItemProvider(
                 CodewhispererInlineCompletionLanguages,
                 this.inlineCompletionProvider
             )
         }
-        commands.registerCommand('aws.sample-vscode-ext-amazonq.accept', onInlineAcceptance)
+        commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance)
 
-        const oninlineRejection = async (sessionId: string, itemId: string) => {
+        const onInlineRejection = async (sessionId: string, itemId: string) => {
             await commands.executeCommand('editor.action.inlineSuggest.hide')
             // TODO: also log the seen state for other suggestions in session
             const params: LogInlineCompletionSessionResultsParams = {
@@ -114,39 +120,49 @@ export class InlineCompletionManager implements Disposable {
                     },
                 },
             }
-            this.languageClient.sendNotification(logInlineCompletionSessionResultsNotificationType as any, params)
+            this.languageClient.sendNotification(this.logSessionResultMessageName, params)
             this.disposable.dispose()
             this.disposable = languages.registerInlineCompletionItemProvider(
                 CodewhispererInlineCompletionLanguages,
                 this.inlineCompletionProvider
             )
         }
-        commands.registerCommand('aws.sample-vscode-ext-amazonq.reject', oninlineRejection)
+        commands.registerCommand('aws.amazonq.rejectInline', onInlineRejection)
 
         /*
             We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item
             To show prev. and next. recommendation we need to re-register a new provider with the previous or next item
         */
 
         const prevCommandHandler = async () => {
-            SessionManager.instance.decrementActiveIndex()
+            this.sessionManager.decrementActiveIndex()
             await commands.executeCommand('editor.action.inlineSuggest.hide')
             this.disposable.dispose()
             this.disposable = languages.registerInlineCompletionItemProvider(
                 CodewhispererInlineCompletionLanguages,
-                new AmazonQInlineCompletionItemProvider(this.languageClient, false)
+                new AmazonQInlineCompletionItemProvider(
+                    this.languageClient,
+                    this.recommendationService,
+                    this.sessionManager,
+                    false
+                )
             )
             await commands.executeCommand('editor.action.inlineSuggest.trigger')
         }
         commands.registerCommand('editor.action.inlineSuggest.showPrevious', prevCommandHandler)
 
         const nextCommandHandler = async () => {
-            SessionManager.instance.incrementActiveIndex()
+            this.sessionManager.incrementActiveIndex()
             await commands.executeCommand('editor.action.inlineSuggest.hide')
             this.disposable.dispose()
             this.disposable = languages.registerInlineCompletionItemProvider(
                 CodewhispererInlineCompletionLanguages,
-                new AmazonQInlineCompletionItemProvider(this.languageClient, false)
+                new AmazonQInlineCompletionItemProvider(
+                    this.languageClient,
+                    this.recommendationService,
+                    this.sessionManager,
+                    false
+                )
             )
             await commands.executeCommand('editor.action.inlineSuggest.trigger')
         }
@@ -157,6 +173,8 @@ export class InlineCompletionManager implements Disposable {
 export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
     constructor(
         private readonly languageClient: LanguageClient,
+        private readonly recommendationService: RecommendationService,
+        private readonly sessionManager: SessionManager,
         private readonly isNewSesion: boolean = true
     ) {}
 
@@ -168,7 +186,7 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
     ): Promise<InlineCompletionItem[] | InlineCompletionList> {
         if (this.isNewSesion) {
             // make service requests if it's a new session
-            await RecommendationService.instance.getAllRecommendations(
+            await this.recommendationService.getAllRecommendations(
                 this.languageClient,
                 document,
                 position,
@@ -177,6 +195,23 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
             )
         }
         // get active item from session for displaying
-        return SessionManager.instance.getActiveRecommendation()
+        const items = this.sessionManager.getActiveRecommendation()
+        const session = this.sessionManager.getActiveSession()
+        if (!session || !items.length) {
+            return []
+        }
+        for (const item of items) {
+            item.command = {
+                command: 'aws.amazonq.acceptInline',
+                title: 'On acceptance',
+                arguments: [
+                    session.sessionId,
+                    item.itemId,
+                    session.requestStartTime,
+                    session.firstCompletionDisplayLatency,
+                ],
+            }
+        }
+        return items as InlineCompletionItem[]
     }
 }

packages/amazonq/src/app/inline/recommendationService.ts

@@ -13,11 +13,7 @@ import { LanguageClient } from 'vscode-languageclient'
 import { SessionManager } from './sessionManager'
 
 export class RecommendationService {
-    static #instance: RecommendationService
-
-    public static get instance() {
-        return (this.#instance ??= new this())
-    }
+    constructor(private readonly sessionManager: SessionManager) {}
 
     async getAllRecommendations(
         languageClient: LanguageClient,
@@ -43,15 +39,20 @@ export class RecommendationService {
         )
 
         const firstCompletionDisplayLatency = Date.now() - requestStartTime
-        SessionManager.instance.startSession(firstResult.sessionId, firstResult.items, firstCompletionDisplayLatency)
+        this.sessionManager.startSession(
+            firstResult.sessionId,
+            firstResult.items,
+            requestStartTime,
+            firstCompletionDisplayLatency
+        )
 
         if (firstResult.partialResultToken) {
             // If there are more results to fetch, handle them in the background
             this.processRemainingRequests(languageClient, request, firstResult, token).catch((error) => {
                 languageClient.warn(`Error when getting suggestions: ${error}`)
             })
         } else {
-            SessionManager.instance.closeSession()
+            this.sessionManager.closeSession()
         }
     }
 
@@ -69,9 +70,9 @@ export class RecommendationService {
                 request,
                 token
             )
-            SessionManager.instance.updateSessionSuggestions(result.items)
+            this.sessionManager.updateSessionSuggestions(result.items)
             nextToken = result.partialResultToken
         }
-        SessionManager.instance.closeSession()
+        this.sessionManager.closeSession()
     }
 }

packages/amazonq/src/app/inline/sessionManager.ts

@@ -4,35 +4,33 @@
  */
 
 import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes-types/inlineCompletionWithReferences'
-import { InlineCompletionItem } from 'vscode'
 
 // TODO: add more needed data to the session interface
 interface CodeWhispererSession {
     sessionId: string
     suggestions: InlineCompletionItemWithReferences[]
     // TODO: might need to convert to enum states
     isRequestInProgress: boolean
+    requestStartTime: number
     firstCompletionDisplayLatency?: number
 }
 
 export class SessionManager {
-    static #instance: SessionManager
     private activeSession?: CodeWhispererSession
     private activeIndex: number = 0
-
-    public static get instance() {
-        return (this.#instance ??= new this())
-    }
+    constructor() {}
 
     public startSession(
         sessionId: string,
         suggestions: InlineCompletionItemWithReferences[],
+        requestStartTime: number,
         firstCompletionDisplayLatency?: number
     ) {
         this.activeSession = {
             sessionId,
             suggestions,
             isRequestInProgress: true,
+            requestStartTime,
             firstCompletionDisplayLatency,
         }
         this.activeIndex = 0
@@ -45,6 +43,10 @@ export class SessionManager {
         this.activeSession.isRequestInProgress = false
     }
 
+    public getActiveSession() {
+        return this.activeSession
+    }
+
     public updateSessionSuggestions(suggestions: InlineCompletionItemWithReferences[]) {
         if (!this.activeSession) {
             return
@@ -69,7 +71,7 @@ export class SessionManager {
         In order to keep track of the right suggestion state, and for features such as reference tracker, this hack is still needed
      */
 
-    public getActiveRecommendation(): InlineCompletionItem[] {
+    public getActiveRecommendation(): InlineCompletionItemWithReferences[] {
         let suggestionCount = this.activeSession?.suggestions.length
         if (!suggestionCount) {
             return []
@@ -78,7 +80,7 @@ export class SessionManager {
             suggestionCount += 1
         }
 
-        const activeSuggestion = this.activeSession?.suggestions[this.activeIndex] as InlineCompletionItem
+        const activeSuggestion = this.activeSession?.suggestions[this.activeIndex]
         if (!activeSuggestion) {
             return []
         }

packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts

@@ -69,16 +69,16 @@ describe('InlineCompletionManager', () => {
         })
 
         it('should register accept and reject commands', () => {
-            assert(registerCommandStub.calledWith('aws.sample-vscode-ext-amazonq.accept'))
-            assert(registerCommandStub.calledWith('aws.sample-vscode-ext-amazonq.reject'))
+            assert(registerCommandStub.calledWith('aws.amazonq.acceptInline'))
+            assert(registerCommandStub.calledWith('aws.amazonq.rejectInline'))
         })
 
         describe('onInlineAcceptance', () => {
             it('should send notification and re-register provider on acceptance', async () => {
                 // Get the acceptance handler
                 const acceptanceHandler = registerCommandStub
                     .getCalls()
-                    ?.find((call) => call.args[0] === 'aws.sample-vscode-ext-amazonq.accept')?.args[1]
+                    ?.find((call) => call.args[0] === 'aws.amazonq.acceptInline')?.args[1]
 
                 const sessionId = 'test-session'
                 const itemId = 'test-item'
@@ -90,7 +90,7 @@ describe('InlineCompletionManager', () => {
                 assert(sendNotificationStub.calledOnce)
                 assert(
                     sendNotificationStub.calledWith(
-                        sinon.match.any,
+                        'aws/logInlineCompletionSessionResults',
                         sinon.match({
                             sessionId,
                             completionSessionResult: {
@@ -114,7 +114,7 @@ describe('InlineCompletionManager', () => {
                 // Get the rejection handler
                 const rejectionHandler = registerCommandStub
                     .getCalls()
-                    .find((call) => call.args[0] === 'aws.sample-vscode-ext-amazonq.reject')?.args[1]
+                    .find((call) => call.args[0] === 'aws.amazonq.rejectInline')?.args[1]
 
                 const sessionId = 'test-session'
                 const itemId = 'test-item'
@@ -125,7 +125,7 @@ describe('InlineCompletionManager', () => {
                 assert(sendNotificationStub.calledOnce)
                 assert(
                     sendNotificationStub.calledWith(
-                        sinon.match.any,
+                        'aws/logInlineCompletionSessionResults',
                         sinon.match({
                             sessionId,
                             completionSessionResult: {

packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts

@@ -27,7 +27,8 @@ describe('RecommendationService', () => {
         insertText: 'ItemTwo',
     } as InlineCompletionItem
     const mockPartialResultToken = 'some-random-token'
-    const service = RecommendationService.instance
+    const sessionManager = new SessionManager()
+    const service = new RecommendationService(sessionManager)
 
     beforeEach(() => {
         sandbox = sinon.createSandbox()
@@ -41,7 +42,7 @@ describe('RecommendationService', () => {
 
     afterEach(() => {
         sandbox.restore()
-        SessionManager.instance.clear()
+        sessionManager.clear()
     })
 
     describe('getAllRecommendations', () => {
@@ -68,7 +69,7 @@ describe('RecommendationService', () => {
             })
 
             // Verify session management
-            const items = SessionManager.instance.getActiveRecommendation()
+            const items = sessionManager.getActiveRecommendation()
             assert.deepStrictEqual(items, [mockInlineCompletionItemOne])
         })
 
@@ -108,10 +109,10 @@ describe('RecommendationService', () => {
             })
 
             // Verify session management
-            const items = SessionManager.instance.getActiveRecommendation()
+            const items = sessionManager.getActiveRecommendation()
             assert.deepStrictEqual(items, [mockInlineCompletionItemOne, { insertText: '1' } as InlineCompletionItem])
-            SessionManager.instance.incrementActiveIndex()
-            const items2 = SessionManager.instance.getActiveRecommendation()
+            sessionManager.incrementActiveIndex()
+            const items2 = sessionManager.getActiveRecommendation()
             assert.deepStrictEqual(items2, [mockInlineCompletionItemTwo, { insertText: '1' } as InlineCompletionItem])
         })
     })