fix(amazonq): fix issue with reject

Author: zixlin7

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

@@ -79,9 +79,19 @@ export class InlineCompletionManager implements Disposable {
         }
         commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance)
 
-        const onInlineRejection = async (sessionId: string, itemId: string) => {
+        const onInlineRejection = async () => {
             await commands.executeCommand('editor.action.inlineSuggest.hide')
             // TODO: also log the seen state for other suggestions in session
+            this.disposable.dispose()
+            this.disposable = languages.registerInlineCompletionItemProvider(
+                CodeWhispererConstants.platformLanguageIds,
+                this.inlineCompletionProvider
+            )
+            const sessionId = this.sessionManager.getActiveSession()?.sessionId
+            const itemId = this.sessionManager.getActiveRecommendation()[0]?.itemId
+            if (!sessionId || !itemId) {
+                return
+            }
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: sessionId,
                 completionSessionResult: {
@@ -93,13 +103,8 @@ export class InlineCompletionManager implements Disposable {
                 },
             }
             this.languageClient.sendNotification(this.logSessionResultMessageName, params)
-            this.disposable.dispose()
-            this.disposable = languages.registerInlineCompletionItemProvider(
-                CodeWhispererConstants.platformLanguageIds,
-                this.inlineCompletionProvider
-            )
         }
-        commands.registerCommand('aws.amazonq.rejectInline', onInlineRejection)
+        commands.registerCommand('aws.amazonq.rejectCodeSuggestion', onInlineRejection)
 
         /*
             We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item

packages/amazonq/src/lsp/activation.ts

@@ -18,9 +18,9 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
             Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
                 await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
             }),
-            Commands.declare('aws.amazonq.rejectCodeSuggestion', () => async () => {
-                await vscode.commands.executeCommand('editor.action.inlineSuggest.hide')
-            }).register()
+            vscode.workspace.onDidCloseTextDocument(async () => {
+                await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
+            })
         )
     } catch (err) {
         const e = err as ToolkitError

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

@@ -17,6 +17,8 @@ describe('InlineCompletionManager', () => {
     let executeCommandStub: sinon.SinonStub
     let disposableStub: sinon.SinonStub
     let sandbox: sinon.SinonSandbox
+    let getActiveSessionStub: sinon.SinonStub
+    let getActiveRecommendationStub: sinon.SinonStub
 
     beforeEach(() => {
         sandbox = sinon.createSandbox()
@@ -38,6 +40,8 @@ describe('InlineCompletionManager', () => {
         } as unknown as LanguageClient
 
         manager = new InlineCompletionManager(languageClient)
+        getActiveSessionStub = sandbox.stub(manager['sessionManager'], 'getActiveSession')
+        getActiveRecommendationStub = sandbox.stub(manager['sessionManager'], 'getActiveRecommendation')
     })
 
     afterEach(() => {
@@ -70,7 +74,7 @@ describe('InlineCompletionManager', () => {
 
         it('should register accept and reject commands', () => {
             assert(registerCommandStub.calledWith('aws.amazonq.acceptInline'))
-            assert(registerCommandStub.calledWith('aws.amazonq.rejectInline'))
+            assert(registerCommandStub.calledWith('aws.amazonq.rejectCodeSuggestion'))
         })
 
         describe('onInlineAcceptance', () => {
@@ -114,12 +118,23 @@ describe('InlineCompletionManager', () => {
                 // Get the rejection handler
                 const rejectionHandler = registerCommandStub
                     .getCalls()
-                    .find((call) => call.args[0] === 'aws.amazonq.rejectInline')?.args[1]
+                    .find((call) => call.args[0] === 'aws.amazonq.rejectCodeSuggestion')?.args[1]
 
                 const sessionId = 'test-session'
                 const itemId = 'test-item'
+                const mockSuggestion = {
+                    itemId,
+                    insertText: 'test',
+                }
 
-                await rejectionHandler(sessionId, itemId)
+                getActiveSessionStub.returns({
+                    sessionId: 'test-session',
+                    suggestions: [mockSuggestion],
+                    isRequestInProgress: false,
+                    requestStartTime: Date.now(),
+                })
+                getActiveRecommendationStub.returns([mockSuggestion])
+                await rejectionHandler()
 
                 assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide'))
                 assert(sendNotificationStub.calledOnce)