Merge branch 'feature/ui-e2e-tests' into inline-test

Author: surajrdy-aws

packages/amazonq/src/app/inline/EditRendering/displayImage.ts

@@ -3,16 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { getLogger, setContext } from 'aws-core-vscode/shared'
+import { getContext, getLogger, setContext } from 'aws-core-vscode/shared'
 import * as vscode from 'vscode'
-import { diffLines } from 'diff'
+import { applyPatch, diffLines } from 'diff'
 import { LanguageClient } from 'vscode-languageclient'
 import { CodeWhispererSession } from '../sessionManager'
 import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-runtimes/protocol'
 import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes/protocol'
 import path from 'path'
 import { imageVerticalOffset } from './svgGenerator'
-import { AmazonQInlineCompletionItemProvider } from '../completion'
+import { EditSuggestionState } from '../editSuggestionState'
+import type { AmazonQInlineCompletionItemProvider } from '../completion'
 import { vsCodeState } from 'aws-core-vscode/codewhisperer'
 
 export class EditDecorationManager {
@@ -136,6 +137,7 @@ export class EditDecorationManager {
         await this.clearDecorations(editor)
 
         await setContext('aws.amazonq.editSuggestionActive' as any, true)
+        EditSuggestionState.setEditSuggestionActive(true)
 
         this.acceptHandler = onAccept
         this.rejectHandler = onReject
@@ -166,6 +168,7 @@ export class EditDecorationManager {
         this.acceptHandler = undefined
         this.rejectHandler = undefined
         await setContext('aws.amazonq.editSuggestionActive' as any, false)
+        EditSuggestionState.setEditSuggestionActive(false)
     }
 
     /**
@@ -270,6 +273,28 @@ function getEndOfEditPosition(originalCode: string, newCode: string): vscode.Pos
     return editor ? editor.selection.active : new vscode.Position(0, 0)
 }
 
+/**
+ * Helper function to create discard telemetry params
+ */
+function createDiscardTelemetryParams(
+    session: CodeWhispererSession,
+    item: InlineCompletionItemWithReferences
+): LogInlineCompletionSessionResultsParams {
+    return {
+        sessionId: session.sessionId,
+        completionSessionResult: {
+            [item.itemId]: {
+                seen: false,
+                accepted: false,
+                discarded: true,
+            },
+        },
+        totalSessionDisplayTime: Date.now() - session.requestStartTime,
+        firstCompletionDisplayLatency: session.firstCompletionDisplayLatency,
+        isInlineEdit: true,
+    }
+}
+
 /**
  * Helper function to display SVG decorations
  */
@@ -286,6 +311,41 @@ export async function displaySvgDecoration(
 ) {
     const originalCode = editor.document.getText()
 
+    // Check if a completion suggestion is currently active - if so, discard edit suggestion
+    if (inlineCompletionProvider && (await inlineCompletionProvider.isCompletionActive())) {
+        // Emit DISCARD telemetry for edit suggestion that can't be shown due to active completion
+        const params = createDiscardTelemetryParams(session, item)
+        languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
+        getLogger().info('Edit suggestion discarded due to active completion suggestion')
+        return
+    }
+
+    const isPatchValid = applyPatch(editor.document.getText(), item.insertText as string)
+    if (!isPatchValid) {
+        const params = createDiscardTelemetryParams(session, item)
+        // TODO: this session is closed on flare side hence discarded is not emitted in flare
+        languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
+        return
+    }
+    const documentChangeListener = vscode.workspace.onDidChangeTextDocument((e) => {
+        if (e.contentChanges.length <= 0) {
+            return
+        }
+        if (e.document !== editor.document) {
+            return
+        }
+        if (vsCodeState.isCodeWhispererEditing) {
+            return
+        }
+        if (getContext('aws.amazonq.editSuggestionActive') === false) {
+            return
+        }
+
+        const isPatchValid = applyPatch(e.document.getText(), item.insertText as string)
+        if (!isPatchValid) {
+            void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
+        }
+    })
     await decorationManager.displayEditSuggestion(
         editor,
         svgImage,
@@ -310,6 +370,7 @@ export async function displaySvgDecoration(
             editor.selection = new vscode.Selection(endPosition, endPosition)
 
             await decorationManager.clearDecorations(editor)
+            documentChangeListener.dispose()
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
@@ -343,6 +404,7 @@ export async function displaySvgDecoration(
             // Handle reject
             getLogger().info('Edit suggestion rejected')
             await decorationManager.clearDecorations(editor)
+            documentChangeListener.dispose()
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {

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

@@ -237,6 +237,53 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
         await vscode.commands.executeCommand(`aws.amazonq.checkInlineSuggestionVisibility`)
     }
 
+    /**
+     * Check if a completion suggestion is currently active/displayed
+     */
+    public async isCompletionActive(): Promise<boolean> {
+        const session = this.sessionManager.getActiveSession()
+        if (session === undefined || !session.displayed || session.suggestions.some((item) => item.isInlineEdit)) {
+            return false
+        }
+
+        // Use VS Code command to check if inline suggestion is actually visible on screen
+        // This command only executes when inlineSuggestionVisible context is true
+        await vscode.commands.executeCommand('aws.amazonq.checkInlineSuggestionVisibility')
+        const isInlineSuggestionVisible = performance.now() - session.lastVisibleTime < 50
+        return isInlineSuggestionVisible
+    }
+
+    /**
+     * Batch discard telemetry for completion suggestions when edit suggestion is active
+     */
+    public batchDiscardTelemetryForEditSuggestion(items: any[], session: any): void {
+        // Emit DISCARD telemetry for completion suggestions that can't be shown due to active edit
+        const completionSessionResult: {
+            [key: string]: { seen: boolean; accepted: boolean; discarded: boolean }
+        } = {}
+
+        for (const item of items) {
+            if (!item.isInlineEdit && item.itemId) {
+                completionSessionResult[item.itemId] = {
+                    seen: false,
+                    accepted: false,
+                    discarded: true,
+                }
+            }
+        }
+
+        // Send single telemetry event for all discarded items
+        if (Object.keys(completionSessionResult).length > 0) {
+            const params: LogInlineCompletionSessionResultsParams = {
+                sessionId: session.sessionId,
+                completionSessionResult,
+                firstCompletionDisplayLatency: session.firstCompletionDisplayLatency,
+                totalSessionDisplayTime: performance.now() - session.requestStartTime,
+            }
+            this.languageClient.sendNotification(this.logSessionResultMessageName, params)
+        }
+    }
+
     // this method is automatically invoked by VS Code as user types
     async provideInlineCompletionItems(
         document: TextDocument,

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

@@ -0,0 +1,19 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Manages the state of edit suggestions to avoid circular dependencies
+ */
+export class EditSuggestionState {
+    private static isEditSuggestionCurrentlyActive = false
+
+    static setEditSuggestionActive(active: boolean): void {
+        this.isEditSuggestionCurrentlyActive = active
+    }
+
+    static isEditSuggestionActive(): boolean {
+        return this.isEditSuggestionCurrentlyActive
+    }
+}

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

@@ -24,6 +24,7 @@ import { getLogger } from 'aws-core-vscode/shared'
 import { DocumentEventListener } from './documentEventListener'
 import { getOpenFilesInWindow } from 'aws-core-vscode/utils'
 import { asyncCallWithTimeout } from '../../util/timeoutUtil'
+import { EditSuggestionState } from './editSuggestionState'
 
 export interface GetAllRecommendationsOptions {
     emitTelemetry?: boolean
@@ -132,7 +133,7 @@ export class RecommendationService {
              * Completions use PartialResultToken with single 1 call of [getAllRecommendations].
              * Edits leverage partialResultToken to achieve EditStreak such that clients can pull all continuous suggestions generated by the model within 1 EOS block.
              */
-            if (!isTriggerByDeletion && !request.partialResultToken) {
+            if (!isTriggerByDeletion && !request.partialResultToken && !EditSuggestionState.isEditSuggestionActive()) {
                 const completionPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
                     inlineCompletionWithReferencesRequestType.method,
                     request,

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

@@ -26,6 +26,8 @@ export interface CodeWhispererSession {
     triggerOnAcceptance?: boolean
     // whether any suggestion in this session was displayed on screen
     displayed: boolean
+    // timestamp when the suggestion was last visible
+    lastVisibleTime: number
 }
 
 export class SessionManager {
@@ -52,6 +54,7 @@ export class SessionManager {
             firstCompletionDisplayLatency,
             diagnosticsBeforeAccept,
             displayed: false,
+            lastVisibleTime: 0,
         }
         this._currentSuggestionIndex = 0
     }
@@ -134,6 +137,7 @@ export class SessionManager {
     public checkInlineSuggestionVisibility() {
         if (this.activeSession) {
             this.activeSession.displayed = true
+            this.activeSession.lastVisibleTime = performance.now()
         }
     }
 

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

@@ -15,6 +15,7 @@ import { CursorUpdateManager } from '../../../../../src/app/inline/cursorUpdateM
 import { CodeWhispererStatusBarManager } from 'aws-core-vscode/codewhisperer'
 import { globals } from 'aws-core-vscode/shared'
 import { DocumentEventListener } from '../../../../../src/app/inline/documentEventListener'
+import { EditSuggestionState } from '../../../../../src/app/inline/editSuggestionState'
 
 const completionApi = 'aws/textDocument/inlineCompletionWithReferences'
 const editApi = 'aws/textDocument/editCompletion'
@@ -325,5 +326,77 @@ describe('RecommendationService', () => {
                 console.error = originalConsoleError
             }
         })
+
+        it('should not make completion request when edit suggestion is active', async () => {
+            // Mock EditSuggestionState to return true (edit suggestion is active)
+            const isEditSuggestionActiveStub = sandbox.stub(EditSuggestionState, 'isEditSuggestionActive').returns(true)
+
+            const mockResult = {
+                sessionId: 'test-session',
+                items: [mockInlineCompletionItemOne],
+                partialResultToken: undefined,
+            }
+
+            sendRequestStub.resolves(mockResult)
+
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true,
+                mockDocumentEventListener
+            )
+
+            // Verify sendRequest was called only for edit API, not completion API
+            const cs = sendRequestStub.getCalls()
+            const completionCalls = cs.filter((c) => c.firstArg === completionApi)
+            const editCalls = cs.filter((c) => c.firstArg === editApi)
+
+            assert.strictEqual(cs.length, 1) // Only edit call
+            assert.strictEqual(completionCalls.length, 0) // No completion calls
+            assert.strictEqual(editCalls.length, 1) // One edit call
+
+            // Verify the stub was called
+            sinon.assert.calledOnce(isEditSuggestionActiveStub)
+        })
+
+        it('should make completion request when edit suggestion is not active', async () => {
+            // Mock EditSuggestionState to return false (no edit suggestion active)
+            const isEditSuggestionActiveStub = sandbox
+                .stub(EditSuggestionState, 'isEditSuggestionActive')
+                .returns(false)
+
+            const mockResult = {
+                sessionId: 'test-session',
+                items: [mockInlineCompletionItemOne],
+                partialResultToken: undefined,
+            }
+
+            sendRequestStub.resolves(mockResult)
+
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true,
+                mockDocumentEventListener
+            )
+
+            // Verify sendRequest was called for both APIs
+            const cs = sendRequestStub.getCalls()
+            const completionCalls = cs.filter((c) => c.firstArg === completionApi)
+            const editCalls = cs.filter((c) => c.firstArg === editApi)
+
+            assert.strictEqual(cs.length, 2) // Both calls
+            assert.strictEqual(completionCalls.length, 1) // One completion call
+            assert.strictEqual(editCalls.length, 1) // One edit call
+
+            // Verify the stub was called
+            sinon.assert.calledOnce(isEditSuggestionActiveStub)
+        })
     })
 })